From 5dc7cdc2eaec25c011aa472e795e3d1d346b0e81 Mon Sep 17 00:00:00 2001 From: Levi Neely Date: Wed, 12 Aug 2026 09:30:39 +0200 Subject: [PATCH] Make Request() always context-aware Remove RequestCtx variant - all Request handlers now take context. This is required for proper cancellation of blocking operations like tool execution when the client disconnects or stop is called. --- cmd/olliesrv/internal/fs/spec.go | 12 +++++------ cmd/toolsrv/internal/server/server.go | 6 +++--- virtfs/builder.go | 31 +++------------------------ virtfs/decl.go | 12 +++-------- 4 files changed, 15 insertions(+), 46 deletions(-) diff --git a/cmd/olliesrv/internal/fs/spec.go b/cmd/olliesrv/internal/fs/spec.go index 0e20217..8387f12 100644 --- a/cmd/olliesrv/internal/fs/spec.go +++ b/cmd/olliesrv/internal/fs/spec.go @@ -109,7 +109,7 @@ func buildTreeSpec(cfg *Config) virtfs.FsNodeDecl { ), virtfs.FileNode("ctl", 0666, virtfs.Doc("Server control. Write: 'invalidate', 'kill'"), - virtfs.Request(func(data []byte) ([]byte, error) { + virtfs.Request(func(_ context.Context, data []byte) ([]byte, error) { return dispatch(map[string]func([]string) ([]byte, error){ "invalidate": func(_ []string) ([]byte, error) { if mc != nil { @@ -140,7 +140,7 @@ func buildTreeSpec(cfg *Config) virtfs.FsNodeDecl { ), virtfs.FileNode("generate", 0666, virtfs.Doc("One-shot LLM generation"), - virtfs.Request(func(data []byte) ([]byte, error) { + virtfs.Request(func(_ context.Context, data []byte) ([]byte, error) { var req backend.GenerateRequest if err := json.Unmarshal(data, &req); err != nil { req.Prompt = strings.TrimSpace(string(data)) @@ -219,7 +219,7 @@ func buildTreeSpec(cfg *Config) virtfs.FsNodeDecl { virtfs.Doc("Session management"), virtfs.FileNode("new", 0666, virtfs.Doc("Create session"), - virtfs.Request(func(data []byte) ([]byte, error) { + virtfs.Request(func(_ context.Context, data []byte) ([]byte, error) { args := strings.Fields(string(data)) name, remote := "", "" for _, arg := range args { @@ -310,7 +310,7 @@ func buildSessionChildren( }), ), virtfs.FileNode("ctl", 0666, - virtfs.Request(func(data []byte) ([]byte, error) { + virtfs.Request(func(_ context.Context, data []byte) ([]byte, error) { return dispatch(map[string]func([]string) ([]byte, error){ "kill": func(_ []string) ([]byte, error) { return []byte("ok\n"), removeFn() }, ".": func(_ []string) ([]byte, error) { return []byte("ok\n"), removeFn() }, @@ -372,7 +372,7 @@ func buildSessionChildren( ), virtfs.DirNode("agent", virtfs.FileNode("new", 0666, - virtfs.Request(func(data []byte) ([]byte, error) { + virtfs.Request(func(_ context.Context, data []byte) ([]byte, error) { ag, err := session.CreateAgent(s.Name(), strings.Fields(string(data))) if err != nil { return nil, fmt.Errorf("create agent: %w", err) @@ -584,7 +584,7 @@ func buildAgentChildren(a *agent.Agent, s *session.Session) []virtfs.FsNodeDecl }), ), virtfs.FileNode("ctl", 0666, - virtfs.Request(func(data []byte) ([]byte, error) { + virtfs.Request(func(_ context.Context, data []byte) ([]byte, error) { return dispatch(map[string]func([]string) ([]byte, error){ "kill": func(_ []string) ([]byte, error) { s.RemoveAgent(a.ID()) diff --git a/cmd/toolsrv/internal/server/server.go b/cmd/toolsrv/internal/server/server.go index 074912f..f99c167 100644 --- a/cmd/toolsrv/internal/server/server.go +++ b/cmd/toolsrv/internal/server/server.go @@ -117,7 +117,7 @@ func Spec(srv *Server) virtfs.FsNodeDecl { ), virtfs.FileNode("tools", 0666, virtfs.Doc("Tools: write agent ID, read tool list (rdwr)"), - virtfs.Request(func(data []byte) ([]byte, error) { + virtfs.Request(func(_ context.Context, data []byte) ([]byte, error) { agentID := strings.TrimSpace(string(data)) if agentID == "" { return nil, fmt.Errorf("agent ID required") @@ -164,7 +164,7 @@ func Spec(srv *Server) virtfs.FsNodeDecl { virtfs.DirNode("proc", virtfs.FileNode("new", 0666, virtfs.Doc("Execute tool: write token + tool + args, read result (blocking)"), - virtfs.RequestCtx(func(ctx context.Context, data []byte) ([]byte, error) { + virtfs.Request(func(ctx context.Context, data []byte) ([]byte, error) { result, err := srv.Fs.HandleProcNew(ctx, string(data)) if err != nil { return nil, err @@ -174,7 +174,7 @@ func Spec(srv *Server) virtfs.FsNodeDecl { ), virtfs.FileNode("new.bg", 0666, virtfs.Doc("Execute tool in background: write token + tool + args, read pid"), - virtfs.RequestCtx(func(ctx context.Context, data []byte) ([]byte, error) { + virtfs.Request(func(ctx context.Context, data []byte) ([]byte, error) { result, err := srv.Fs.HandleProcNewBg(ctx, string(data)) if err != nil { return nil, err diff --git a/virtfs/builder.go b/virtfs/builder.go index e788618..449cad4 100644 --- a/virtfs/builder.go +++ b/virtfs/builder.go @@ -18,7 +18,7 @@ func BuildTree(spec FsNodeDecl) *Tree { func validate(d *FsNodeDecl, path string) { hasChildren := len(d.Children) > 0 hasBindings := d.Bindings != nil - hasFileHandlers := d.Read != nil || d.Write != nil || d.BlockOnce != nil || d.Stream != nil || d.Request != nil || d.RequestCtx != nil + hasFileHandlers := d.Read != nil || d.Write != nil || d.BlockOnce != nil || d.Stream != nil || d.Request != nil isDir := hasChildren @@ -42,11 +42,8 @@ func validate(d *FsNodeDecl, path string) { if d.Request != nil { blocking++ } - if d.RequestCtx != nil { - blocking++ - } if blocking > 1 { - panic(fmt.Sprintf("BuildTree: %q: at most one of BlockOnce, Stream, Request, RequestCtx may be set", fullPath)) + panic(fmt.Sprintf("BuildTree: %q: at most one of BlockOnce, Stream, Request may be set", fullPath)) } // Template nodes must have Bindings set. @@ -189,7 +186,7 @@ func statDir(d *FsNodeDecl, uid, gid, name string) (os.FileInfo, error) { } cu := resolveUID(child, uid) cg := resolveGID(child, gid) - hasFileHandlers := child.Read != nil || child.Write != nil || child.BlockOnce != nil || child.Stream != nil || child.Request != nil || child.RequestCtx != nil + hasFileHandlers := child.Read != nil || child.Write != nil || child.BlockOnce != nil || child.Stream != nil || child.Request != nil isChildDir := len(child.Children) > 0 || (child.Bindings != nil && !hasFileHandlers) if isChildDir { return &SyntheticFileInfo{Name_: name, Mode_: cm | os.ModeDir, IsDir_: true, UID_: cu, GID_: cg}, nil @@ -261,28 +258,6 @@ func openLeaf(d *FsNodeDecl, uid, gid string, mode os.FileMode, name string) (Fi case d.Request != nil: fn := d.Request var result []byte - blockFn = notBlocking - readFn = func() ([]byte, error) { - if result != nil { - return result, nil - } - if d.Read != nil { - return d.Read() - } - return nil, nil - } - writeFn = func(data []byte) error { - r, err := fn(data) - if err != nil { - return err - } - result = r - return nil - } - isRequest = true - case d.RequestCtx != nil: - fn := d.RequestCtx - var result []byte var reqCtx context.Context var reqCancel context.CancelFunc blockFn = notBlocking diff --git a/virtfs/decl.go b/virtfs/decl.go index 2737990..eb5d9fe 100644 --- a/virtfs/decl.go +++ b/virtfs/decl.go @@ -64,8 +64,7 @@ type FsNodeDecl struct { Stream func(ctx context.Context, base string) ([]byte, string, error) // Request — rdwr: write triggers result, read returns it once. - Request func(data []byte) ([]byte, error) - RequestCtx func(ctx context.Context, data []byte) ([]byte, error) // context-aware variant + Request func(ctx context.Context, data []byte) ([]byte, error) // Tremove + Trename. Remove func() error @@ -133,16 +132,11 @@ func BlockOnce(fn func(context.Context, string) ([]byte, string, error)) NodeOpt return func(d *FsNodeDecl) { d.BlockOnce = fn } } -// Request sets a request-response handler (rdwr pattern). -func Request(fn func([]byte) ([]byte, error)) NodeOption { +// Request sets a context-aware request-response handler (rdwr pattern). +func Request(fn func(context.Context, []byte) ([]byte, error)) NodeOption { return func(d *FsNodeDecl) { d.Request = fn } } -// RequestCtx sets a context-aware request-response handler. -func RequestCtx(fn func(context.Context, []byte) ([]byte, error)) NodeOption { - return func(d *FsNodeDecl) { d.RequestCtx = fn } -} - // RemoveNode sets the removal handler for a node. func RemoveNode(fn func() error) NodeOption { return func(d *FsNodeDecl) { d.Remove = fn }