lkn
/
ollie
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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.
This commit is contained in:
Levi Neely 2026-08-12 09:30:39 +02:00
parent 861cb47d13
commit 5dc7cdc2ea
4 changed files with 15 additions and 46 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
cmd/olliesrv/internal/fs/spec.go
View File

@ -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())

6
cmd/toolsrv/internal/server/server.go
View File

@ -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

31
virtfs/builder.go
View File

@ -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

12
virtfs/decl.go
View File

@ -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 }