From d59f49ee54edc375b088b78ddeba0c46e60e5b78 Mon Sep 17 00:00:00 2001 From: Levi Neely Date: Wed, 12 Aug 2026 10:00:44 +0200 Subject: [PATCH] feat: context cancellation for tool calls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire context cancellation through the toolsrv 9P server so that canceling a CallTool context properly terminates the running tool. Changes: - virtfs: Add Close() to File interface with CloseFn for Request handlers - toolsrv/p9.go: Handle requests concurrently to allow Tclunk during blocking Twrite; call entry.Close() in handleClunk - toolsrv/internal/server/proc.go: Check ctx.Done() while waiting for tool completion in NewProc - Integration tests for context cancellation chain The cancellation chain: client cancels ctx → fid.Close() → Tclunk → handleClunk calls entry.Close() → closeFn cancels reqCtx → NewProc sees ctx.Done() and returns ctx.Err() --- cmd/toolsrv/integration_test.go | 146 ++++++++++++++++++++++++++++ cmd/toolsrv/internal/server/proc.go | 19 +++- cmd/toolsrv/p9.go | 40 ++++++-- virtfs/builder.go | 23 ++++- virtfs/decl_test.go | 57 ++++++++++- virtfs/file.go | 8 ++ virtfs/tree_test.go | 2 +- 7 files changed, 281 insertions(+), 14 deletions(-) diff --git a/cmd/toolsrv/integration_test.go b/cmd/toolsrv/integration_test.go index b40f5d8..e63b0b4 100644 --- a/cmd/toolsrv/integration_test.go +++ b/cmd/toolsrv/integration_test.go @@ -60,6 +60,11 @@ func startTestServer(t *testing.T) (socketPath string, cleanup func()) { t.Fatalf("failed to create config dir: %v", err) } sandboxCfg := `filesystem: + ro: + - /proc + - /sys + rw: + - /dev rox: - /usr - /lib @@ -79,6 +84,15 @@ network: t.Fatalf("failed to write sandbox.yaml: %v", err) } + // Symlink the real tools directory so tests can find them + realToolsPath := toolsrv.ToolsPath() + if _, err := os.Stat(realToolsPath); err == nil { + if err := os.Symlink(realToolsPath, filepath.Join(cfgDir, "tools")); err != nil { + os.RemoveAll(tmpDir) + t.Fatalf("failed to symlink tools dir: %v", err) + } + } + socketPath = filepath.Join(tmpDir, "toolsrv.sock") cwd, _ := os.Getwd() @@ -503,3 +517,135 @@ func TestIntegration_ProcessKeeperRespawn(t *testing.T) { t.Logf("Successfully verified: old_token=%s new_token=%s", token1[:8], token2[:8]) _ = cwd // silence unused warning } + +// TestIntegration_ToolCancellation tests that canceling the context during a tool +// call causes the call to return quickly (not wait for the full tool execution). +func TestIntegration_ToolCancellation(t *testing.T) { + // Skip if tools directory doesn't exist (CI environment) + if _, err := os.ReadDir(toolsrv.ToolsPath()); err != nil { + t.Skipf("skipping: tools directory not available: %v", err) + } + + socketPath, cleanup := startTestServer(t) + defer cleanup() + + conn, err := toolsrv.Dial(socketPath, "test-secret-cancel") + if err != nil { + t.Fatalf("Dial failed: %v", err) + } + defer conn.Close() + conn.SetAgentID("test-agent-cancel") + + // Load the shell tool + if err := conn.LoadTool("shell"); err != nil { + t.Skipf("skipping: LoadTool failed (shell tool not available): %v", err) + } + + // Create a context that we'll cancel after the tool starts + ctx, cancel := context.WithCancel(context.Background()) + + // Run a long-running command + args := json.RawMessage(`{"cmd": "sleep 30"}`) + + done := make(chan struct{}) + var callErr error + var result json.RawMessage + start := time.Now() + + go func() { + defer close(done) + result, callErr = conn.CallTool(ctx, "shell", args) + }() + + // Wait a bit for the tool to start, then cancel + time.Sleep(200 * time.Millisecond) + cancel() + + // Wait for CallTool to return + select { + case <-done: + // Expected + case <-time.After(5 * time.Second): + t.Fatal("CallTool did not return after context cancellation") + } + + elapsed := time.Since(start) + + // The call should return quickly (not wait 30 seconds) + // This is the key assertion - cancellation caused early return + if elapsed > 3*time.Second { + t.Errorf("CallTool took %v, expected quick return after cancellation", elapsed) + } + + // Note: The error may be nil if the fid was closed before the read completed + // The important thing is that it returned quickly + t.Logf("Context cancellation worked: elapsed=%v, err=%v, result_len=%d", elapsed, callErr, len(result)) +} + +// TestIntegration_ToolCancellation_ContextPropagation tests that context cancellation +// properly propagates through the entire chain: client -> toolsrv -> tool execution. +// This verifies that the client returns quickly when canceled, even if the tool would +// take a long time to complete. +func TestIntegration_ToolCancellation_ContextPropagation(t *testing.T) { + // Skip if tools directory doesn't exist (CI environment) + if _, err := os.ReadDir(toolsrv.ToolsPath()); err != nil { + t.Skipf("skipping: tools directory not available: %v", err) + } + + socketPath, cleanup := startTestServer(t) + defer cleanup() + + conn, err := toolsrv.Dial(socketPath, "test-secret-cancel-prop") + if err != nil { + t.Fatalf("Dial failed: %v", err) + } + defer conn.Close() + conn.SetAgentID("test-agent-cancel-prop") + + // Load the shell tool + if err := conn.LoadTool("shell"); err != nil { + t.Skipf("skipping: LoadTool failed (shell tool not available): %v", err) + } + + // Command: echo immediately, then sleep for a long time + // We verify that cancellation works by checking the call returns quickly + args := json.RawMessage(`{"cmd": "echo started && sleep 30 && echo finished"}`) + + ctx, cancel := context.WithCancel(context.Background()) + + done := make(chan struct{}) + var callErr error + var result json.RawMessage + start := time.Now() + + go func() { + defer close(done) + result, callErr = conn.CallTool(ctx, "shell", args) + }() + + // Wait a bit for the tool to start, then cancel + time.Sleep(300 * time.Millisecond) + cancel() + + // Wait for CallTool to return + select { + case <-done: + t.Logf("CallTool returned: err=%v result=%s", callErr, string(result)) + case <-time.After(5 * time.Second): + t.Fatal("CallTool did not return after cancellation") + } + + elapsed := time.Since(start) + + // The call should return quickly (not wait 30 seconds) + if elapsed > 3*time.Second { + t.Errorf("CallTool took %v, expected quick return after cancellation", elapsed) + } + + // If the result contains "finished", the process wasn't killed + if result != nil && strings.Contains(string(result), "finished") { + t.Error("result contains 'finished' - process was not killed on cancellation") + } + + t.Logf("Context propagation test passed: elapsed=%v", elapsed) +} diff --git a/cmd/toolsrv/internal/server/proc.go b/cmd/toolsrv/internal/server/proc.go index 133431c..bdda58d 100644 --- a/cmd/toolsrv/internal/server/proc.go +++ b/cmd/toolsrv/internal/server/proc.go @@ -285,8 +285,23 @@ func (st *State) NewProc(ctx context.Context, payload string, background bool) ( return fmt.Sprintf("%d", proc.ID), proc.ID, nil } - // Wait for completion - <-proc.done + // Wait for completion or context cancellation + select { + case <-proc.done: + // Normal completion + case <-ctx.Done(): + // Context was canceled - cancel the proc and return error + cancel() + // Wait briefly for the process to clean up + select { + case <-proc.done: + case <-time.After(100 * time.Millisecond): + } + st.procMu.Lock() + delete(st.procs, proc.ID) + st.procMu.Unlock() + return "", proc.ID, ctx.Err() + } proc.mu.Lock() result = proc.output.String() diff --git a/cmd/toolsrv/p9.go b/cmd/toolsrv/p9.go index 12db021..b7c616e 100644 --- a/cmd/toolsrv/p9.go +++ b/cmd/toolsrv/p9.go @@ -58,9 +58,24 @@ func serve9P(ctx context.Context, conn net.Conn, tree *virtfs.Tree, srv *server. srv: srv, } + responses := make(chan *plan9.Fcall, 16) + var wg sync.WaitGroup + + // Writer goroutine + go func() { + for resp := range responses { + if err := plan9.WriteFcall(conn, resp); err != nil { + logger.Warn("write error: %v", err) + return + } + } + }() + for { select { case <-ctx.Done(): + wg.Wait() + close(responses) return default: } @@ -70,14 +85,18 @@ func serve9P(ctx context.Context, conn net.Conn, tree *virtfs.Tree, srv *server. if err != io.EOF { logger.Warn("read error: %v", err) } + wg.Wait() + close(responses) return } - resp := handleFcall(cs, fc, tree) - if err := plan9.WriteFcall(conn, resp); err != nil { - logger.Warn("write error: %v", err) - return - } + // Handle requests concurrently to allow Tclunk during blocking operations + wg.Add(1) + go func(fc *plan9.Fcall) { + defer wg.Done() + resp := handleFcall(cs, fc, tree) + responses <- resp + }(fc) } } @@ -459,9 +478,18 @@ func handleWrite(cs *connState, fc *plan9.Fcall, tree *virtfs.Tree) *plan9.Fcall func handleClunk(cs *connState, fc *plan9.Fcall) *plan9.Fcall { cs.mu.Lock() - delete(cs.fids, fc.Fid) + f, ok := cs.fids[fc.Fid] + if ok { + delete(cs.fids, fc.Fid) + } delete(cs.authFids, fc.Fid) cs.mu.Unlock() + + // Close the file handle to cancel any pending Request handlers. + if ok && f.entry != nil { + f.entry.Close() + } + return &plan9.Fcall{Type: plan9.Rclunk, Tag: fc.Tag} } diff --git a/virtfs/builder.go b/virtfs/builder.go index 449cad4..b35e80e 100644 --- a/virtfs/builder.go +++ b/virtfs/builder.go @@ -279,8 +279,29 @@ func openLeaf(d *FsNodeDecl, uid, gid string, mode os.FileMode, name string) (Fi result = r return nil } - _ = reqCancel // TODO: wire up cancellation on fid close + closeFn := func() error { + if reqCancel != nil { + reqCancel() + } + return nil + } isRequest = true + return &fileConfig{ + StatFn: func() (os.FileInfo, error) { + if d.Stat != nil { + return d.Stat(), nil + } + data, _ := readFn() + return &SyntheticFileInfo{Name_: name, Mode_: mode, Size_: int64(len(data)), UID_: uid, GID_: gid}, nil + }, + ReadFn: readFn, + WriteFn: writeFn, + CloseFn: closeFn, + BlockingReadFn: blockFn, + BlockingRead_: isBlocking, + StreamMode_: isStreaming, + RequestResp_: isRequest, + }, nil default: blockFn = notBlocking } diff --git a/virtfs/decl_test.go b/virtfs/decl_test.go index 1bf7473..71647f5 100644 --- a/virtfs/decl_test.go +++ b/virtfs/decl_test.go @@ -5,6 +5,7 @@ import ( "errors" "os" "testing" + "time" ) // ══════════════════════════════════════════════════════════════════ @@ -31,7 +32,7 @@ func TestNodeOptions(t *testing.T) { }) t.Run("Request", func(t *testing.T) { - spec := FileNode("test", 0666, Request(func(data []byte) ([]byte, error) { + spec := FileNode("test", 0666, Request(func(_ context.Context, data []byte) ([]byte, error) { return nil, nil })) if spec.Request == nil { @@ -280,7 +281,7 @@ func TestOpenLeafStream(t *testing.T) { func TestOpenLeafRequest(t *testing.T) { spec := DirNode("/", - FileNode("request", 0666, Request(func(data []byte) ([]byte, error) { + FileNode("request", 0666, Request(func(_ context.Context, data []byte) ([]byte, error) { return append([]byte("echo:"), data...), nil })), ) @@ -306,7 +307,7 @@ func TestOpenLeafRequest(t *testing.T) { func TestOpenLeafRequestError(t *testing.T) { spec := DirNode("/", - FileNode("request", 0666, Request(func(data []byte) ([]byte, error) { + FileNode("request", 0666, Request(func(_ context.Context, data []byte) ([]byte, error) { return nil, errors.New("request error") })), ) @@ -324,7 +325,7 @@ func TestOpenLeafRequestWithFallbackRead(t *testing.T) { spec := DirNode("/", FileNode("request", 0666, Read(func() ([]byte, error) { return []byte("fallback"), nil }), - Request(func(data []byte) ([]byte, error) { return []byte("response"), nil }), + Request(func(_ context.Context, data []byte) ([]byte, error) { return []byte("response"), nil }), ), ) tree := BuildTree(spec) @@ -341,6 +342,54 @@ func TestOpenLeafRequestWithFallbackRead(t *testing.T) { } } +func TestOpenLeafRequestCancelOnClose(t *testing.T) { + canceled := make(chan struct{}) + started := make(chan struct{}) + spec := DirNode("/", + FileNode("request", 0666, Request(func(ctx context.Context, data []byte) ([]byte, error) { + close(started) + <-ctx.Done() + close(canceled) + return nil, ctx.Err() + })), + ) + tree := BuildTree(spec) + f, err := tree.Open("request") + if err != nil { + t.Fatalf("Open: %v", err) + } + + // Start the request in a goroutine + done := make(chan error, 1) + go func() { + done <- f.Write([]byte("test")) + }() + + // Wait for request to start + <-started + + // Close should cancel the context + f.Close() + + // Verify context was canceled + select { + case <-canceled: + // Success + case <-time.After(time.Second): + t.Fatal("context was not canceled on close") + } + + // Write should return with context canceled error + select { + case err := <-done: + if err == nil || err != context.Canceled { + t.Errorf("Write error = %v, want context.Canceled", err) + } + case <-time.After(time.Second): + t.Fatal("Write did not return after close") + } +} + func TestOpenLeafNoWriteHandler(t *testing.T) { spec := DirNode("/", FileNode("readonly", 0444, Read(func() ([]byte, error) { return []byte("data"), nil })), diff --git a/virtfs/file.go b/virtfs/file.go index f9ea676..912f3d3 100644 --- a/virtfs/file.go +++ b/virtfs/file.go @@ -19,6 +19,7 @@ type File interface { Stat() (os.FileInfo, error) Read() ([]byte, error) Write(data []byte) error + Close() error // called when fid is clunked/closed BlockingRead(ctx context.Context, base string) (content []byte, nextBase string, err error) // Read-side contracts (mutually exclusive — at most one true): @@ -35,6 +36,7 @@ type fileConfig struct { StatFn func() (os.FileInfo, error) ReadFn func() ([]byte, error) WriteFn func([]byte) error + CloseFn func() error BlockingReadFn func(context.Context, string) ([]byte, string, error) BlockingRead_ bool // one-shot waiter: block 5s, return once, EOF on 2nd read StreamMode_ bool // indefinite stream: block forever, keep FD open @@ -44,6 +46,12 @@ type fileConfig struct { func (e *fileConfig) Stat() (os.FileInfo, error) { return e.StatFn() } func (e *fileConfig) Read() ([]byte, error) { return e.ReadFn() } func (e *fileConfig) Write(data []byte) error { return e.WriteFn(data) } +func (e *fileConfig) Close() error { + if e.CloseFn != nil { + return e.CloseFn() + } + return nil +} func (e *fileConfig) BlockingRead(ctx context.Context, base string) ([]byte, string, error) { return e.BlockingReadFn(ctx, base) } diff --git a/virtfs/tree_test.go b/virtfs/tree_test.go index 5d1af4c..254759d 100644 --- a/virtfs/tree_test.go +++ b/virtfs/tree_test.go @@ -256,7 +256,7 @@ func TestGenerateHelpModes(t *testing.T) { spec := DirNode("/", FileNode("ro", 0444, Doc("read only file")), FileNode("rw", 0644, Doc("read write file"), Write(func(_ []byte) error { return nil })), - FileNode("ctl", 0200, Doc("control file"), Request(func(_ []byte) ([]byte, error) { return nil, nil })), + FileNode("ctl", 0200, Doc("control file"), Request(func(_ context.Context, _ []byte) ([]byte, error) { return nil, nil })), ) help := GenerateHelp(spec) if help == "" {