diff --git a/cmd/librecode/chat.go b/cmd/librecode/chat.go index fbd09c47..eb22e28a 100644 --- a/cmd/librecode/chat.go +++ b/cmd/librecode/chat.go @@ -58,6 +58,7 @@ func runChat(cmd *cobra.Command, options chatRunOptions) error { authStorage := container.AuthService().Storage extensionManager := container.ExtensionService().Manager cfg := container.ConfigService().Get() + chatWorkflows := container.ChatWorkflowService() cwd, err := assistant.DefaultCWD("") if err != nil { @@ -75,6 +76,7 @@ func runChat(cmd *cobra.Command, options chatRunOptions) error { Extensions: extensionManager, Resources: &resources, Runtime: runtime, + Workflows: chatWorkflows.Runs, Settings: databaseService.Documents, Models: modelRegistry, Auth: authStorage, diff --git a/cmd/librecode/main.go b/cmd/librecode/main.go index 1510961c..ff4530c1 100644 --- a/cmd/librecode/main.go +++ b/cmd/librecode/main.go @@ -7,9 +7,23 @@ import ( "os" "os/signal" "syscall" + + "github.com/omarluq/librecode/internal/executeworker" ) func main() { + if len(os.Args) == 2 && os.Args[1] == "__execute-worker" { + if err := executeworker.Serve(os.Stdin, os.Stdout); err != nil { + if _, writeErr := fmt.Fprintln(os.Stderr, err); writeErr != nil { + os.Exit(1) + } + + os.Exit(1) + } + + os.Exit(0) + } + os.Exit(run()) } diff --git a/cmd/librecode/prompt.go b/cmd/librecode/prompt.go index d3024421..4ab252c8 100644 --- a/cmd/librecode/prompt.go +++ b/cmd/librecode/prompt.go @@ -1,9 +1,12 @@ package main import ( + "encoding/json" "errors" "fmt" + "os" "strings" + "time" "github.com/samber/oops" "github.com/spf13/cobra" @@ -13,12 +16,17 @@ import ( "github.com/omarluq/librecode/internal/limitio" ) -const promptStdinLimitBytes int64 = 1 << 20 +const ( + promptStdinLimitBytes int64 = 1 << 20 + promptMetricsMode = 0o600 +) type promptRunOptions struct { - SessionID string - SessionName string - Resume bool + SessionID string + SessionName string + ToolStrategy string + MetricsJSON string + Resume bool } func newPromptCmd() *cobra.Command { @@ -35,6 +43,9 @@ func newPromptCmd() *cobra.Command { cmd.Flags().StringVar(&options.SessionID, "session", "", "session id to append to") cmd.Flags().StringVar(&options.SessionName, "name", "", "create a named session") + cmd.Flags().StringVar(&options.ToolStrategy, "tool-strategy", string(assistant.ToolStrategyHybrid), + "tool strategy: hybrid or direct") + cmd.Flags().StringVar(&options.MetricsJSON, "metrics-json", "", "write prompt metrics as JSON") cmd.Flags().BoolVar(&options.Resume, "resume", false, "resume the latest session for this working directory") return cmd @@ -56,6 +67,15 @@ func runPrompt(cmd *cobra.Command, args []string, options promptRunOptions) erro } func validatePromptRunOptions(options promptRunOptions) error { + strategy := assistant.ToolStrategy(options.ToolStrategy) + if strategy == "" { + strategy = assistant.ToolStrategyHybrid + } + + if strategy != assistant.ToolStrategyHybrid && strategy != assistant.ToolStrategyDirect { + return fmt.Errorf("invalid --tool-strategy %q: use hybrid or direct", options.ToolStrategy) + } + if options.Resume && options.SessionID != "" { return errors.New("--resume cannot be used with --session") } @@ -80,9 +100,27 @@ func runPromptWithContainer( return cliError(err, cliResolveWorkingDirectory) } - response, err := runtime.Prompt(cmd.Context(), buildPromptRequest(cwd, message, options)) - if err != nil { - return cliError(err, "run prompt") + strategy := normalizedToolStrategy(options.ToolStrategy) + metrics := new(assistant.RunMetrics) + ctx := assistant.WithToolStrategy(cmd.Context(), strategy) + ctx = assistant.WithRunMetrics(ctx, metrics) + request := buildPromptRequest(cwd, message, options) + request.OnEvent = metrics.ObserveStreamEvent + started := time.Now() + response, promptErr := runtime.Prompt(ctx, request) + elapsed := time.Since(started) + + snapshot := metrics.Snapshot() + measured := &promptMetrics{ + Strategy: string(strategy), Error: errorText(promptErr), + ProviderRoundTrips: snapshot.ProviderRoundTrips, ElapsedMilliseconds: elapsed.Milliseconds(), + InputTokens: snapshot.InputTokens, OutputTokens: snapshot.OutputTokens, + ToolCalls: snapshot.ToolCalls, NestedToolCalls: snapshot.NestedToolCalls, + TraceComplete: snapshot.TraceComplete, Success: promptErr == nil, + } + + if err := promptExecutionError(promptErr, writePromptMetrics(options.MetricsJSON, measured)); err != nil { + return err } if _, err := fmt.Fprintln(cmd.OutOrStdout(), response.Text); err != nil { @@ -92,6 +130,67 @@ func runPromptWithContainer( return nil } +func promptExecutionError(promptErr, metricsErr error) error { + if promptErr == nil { + return metricsErr + } + + promptErr = cliError(promptErr, "run prompt") + if metricsErr != nil { + return errors.Join(promptErr, metricsErr) + } + + return promptErr +} + +type promptMetrics struct { + Strategy string `json:"strategy"` + Error string `json:"error"` + ProviderRoundTrips int `json:"provider_round_trips"` + ElapsedMilliseconds int64 `json:"elapsed_ms"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + ToolCalls int `json:"tool_calls"` + NestedToolCalls int `json:"nested_tool_calls"` + TraceComplete bool `json:"trace_complete"` + Success bool `json:"success"` +} + +func normalizedToolStrategy(value string) assistant.ToolStrategy { + strategy := assistant.ToolStrategy(value) + if strategy == "" { + return assistant.ToolStrategyHybrid + } + + return strategy +} + +func errorText(err error) string { + if err == nil { + return "" + } + + return err.Error() +} + +func writePromptMetrics(path string, metrics *promptMetrics) error { + if strings.TrimSpace(path) == "" { + return nil + } + + encoded, err := json.MarshalIndent(metrics, "", " ") + if err != nil { + return oops.In("cli").Code("encode_prompt_metrics").Wrapf(err, "encode prompt metrics") + } + + encoded = append(encoded, '\n') + if err := os.WriteFile(path, encoded, promptMetricsMode); err != nil { + return oops.In("cli").Code("write_prompt_metrics").Wrapf(err, "write prompt metrics") + } + + return nil +} + func buildPromptRequest(cwd, message string, options promptRunOptions) *assistant.PromptRequest { return &assistant.PromptRequest{ OnEvent: nil, diff --git a/cmd/librecode/prompt_internal_test.go b/cmd/librecode/prompt_internal_test.go index 26ee83d9..1e1062ee 100644 --- a/cmd/librecode/prompt_internal_test.go +++ b/cmd/librecode/prompt_internal_test.go @@ -1,7 +1,10 @@ package main import ( + "encoding/json" "errors" + "os" + "path/filepath" "strings" "testing" @@ -23,30 +26,43 @@ func TestValidatePromptRunOptions(t *testing.T) { options promptRunOptions }{ { - options: promptRunOptions{SessionID: "", SessionName: "", Resume: false}, + options: promptRunOptions{SessionID: "", SessionName: "", ToolStrategy: "", MetricsJSON: "", Resume: false}, name: "default", wantErr: "", }, { - options: promptRunOptions{SessionID: testSessionID, SessionName: "", Resume: false}, + options: promptRunOptions{ + SessionID: testSessionID, SessionName: "", ToolStrategy: "", MetricsJSON: "", Resume: false, + }, name: "session only", wantErr: "", }, { - options: promptRunOptions{SessionID: "", SessionName: "", Resume: true}, + options: promptRunOptions{SessionID: "", SessionName: "", ToolStrategy: "", MetricsJSON: "", Resume: true}, name: "resume only", wantErr: "", }, { - options: promptRunOptions{SessionID: testSessionID, SessionName: "", Resume: true}, + options: promptRunOptions{ + SessionID: testSessionID, SessionName: "", ToolStrategy: "", MetricsJSON: "", Resume: true, + }, name: "resume with session", wantErr: "--resume cannot be used with --session", }, { - options: promptRunOptions{SessionID: "", SessionName: testSessionName, Resume: true}, + options: promptRunOptions{ + SessionID: "", SessionName: testSessionName, ToolStrategy: "hybrid", MetricsJSON: "", Resume: true, + }, name: "resume with name", wantErr: "--resume cannot be used with --name", }, + { + options: promptRunOptions{ + SessionID: "", SessionName: "", ToolStrategy: "invalid", MetricsJSON: "", Resume: false, + }, + name: "invalid strategy", + wantErr: "invalid --tool-strategy", + }, } for _, testCase := range tests { @@ -165,9 +181,7 @@ func TestBuildPromptRequest(t *testing.T) { t.Parallel() request := buildPromptRequest("/workspace", "hello", promptRunOptions{ - SessionID: "session-1", - SessionName: testSessionName, - Resume: true, + SessionID: "session-1", SessionName: testSessionName, ToolStrategy: "", MetricsJSON: "", Resume: true, }) require.NotNil(t, request) @@ -183,6 +197,32 @@ func TestBuildPromptRequest(t *testing.T) { assert.False(t, request.HideUserPrompt) } +func TestWritePromptMetrics(t *testing.T) { + t.Parallel() + + path := filepath.Join(t.TempDir(), "metrics.json") + metrics := &promptMetrics{ + Strategy: "direct", Error: "", ProviderRoundTrips: 2, ElapsedMilliseconds: 12, + InputTokens: 10, OutputTokens: 4, ToolCalls: 1, NestedToolCalls: 0, + TraceComplete: true, Success: true, + } + require.NoError(t, writePromptMetrics(path, metrics)) + + encoded, err := os.ReadFile(filepath.Clean(path)) + require.NoError(t, err) + + var decoded promptMetrics + require.NoError(t, json.Unmarshal(encoded, &decoded)) + assert.Equal(t, *metrics, decoded) +} + +func TestNormalizedToolStrategyDefaultsToHybrid(t *testing.T) { + t.Parallel() + + assert.Equal(t, "hybrid", string(normalizedToolStrategy(""))) + assert.Equal(t, "direct", string(normalizedToolStrategy("direct"))) +} + func TestPromptMessageRejectsStdinAboveLimit(t *testing.T) { t.Parallel() diff --git a/go.mod b/go.mod index c7ace78a..f39d284d 100644 --- a/go.mod +++ b/go.mod @@ -24,6 +24,7 @@ require ( github.com/gofrs/uuid/v5 v5.4.0 github.com/golang-jwt/jwt/v5 v5.3.1 github.com/invopop/jsonschema v0.14.0 + github.com/mvm-sh/mvm v0.5.0 github.com/odvcencio/gotreesitter v0.38.0 github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c github.com/pressly/goose/v3 v3.27.2 diff --git a/go.sum b/go.sum index c1f4b0f5..4f9446bf 100644 --- a/go.sum +++ b/go.sum @@ -135,12 +135,10 @@ github.com/mfridman/interpolate v0.0.2 h1:pnuTK7MQIxxFz1Gr+rjSIx9u7qVjf5VOoM/u6B github.com/mfridman/interpolate v0.0.2/go.mod h1:p+7uk6oE07mpE/Ik1b8EckO0O4ZXiGAfshKBWLUM9Xg= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= +github.com/mvm-sh/mvm v0.5.0 h1:XWII2Y8RLEzvFMuBWSacg6KLUTUp2weBin4lUlFrV+Y= +github.com/mvm-sh/mvm v0.5.0/go.mod h1:2h9+ibS1DzdkMRNjGVs+pkU0blZZXDrEIUpcI93p4XA= github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= -github.com/odvcencio/gotreesitter v0.36.0 h1:WN9pMlzIEOcDoeWDSSeNHC9zpPi7kcZ66sYmf7eknDY= -github.com/odvcencio/gotreesitter v0.36.0/go.mod h1:hBVkghd0paaYAVwd2087vfwdeU984bQbMo9LvpE0moo= -github.com/odvcencio/gotreesitter v0.37.0 h1:9LO82hxtUJKIzDSfqhCZggCYTaxdX4JwXqydbbQi+iE= -github.com/odvcencio/gotreesitter v0.37.0/go.mod h1:hBVkghd0paaYAVwd2087vfwdeU984bQbMo9LvpE0moo= github.com/odvcencio/gotreesitter v0.38.0 h1:6UkjeVOeihIqSilBw3ezUH6UK9korS5TfCmhXF9sFHc= github.com/odvcencio/gotreesitter v0.38.0/go.mod h1:hBVkghd0paaYAVwd2087vfwdeU984bQbMo9LvpE0moo= github.com/oklog/ulid/v2 v2.1.1 h1:suPZ4ARWLOJLegGFiZZ1dFAkqzhMjL3J1TzI+5wHz8s= @@ -357,20 +355,20 @@ lukechampine.com/uint128 v1.1.1/go.mod h1:c4eWIwlEGaxC/+H1VguhU4PHXNWDCDMUlWdIWl lukechampine.com/uint128 v1.2.0/go.mod h1:c4eWIwlEGaxC/+H1VguhU4PHXNWDCDMUlWdIWl2j1gk= modernc.org/cc/v3 v3.37.0/go.mod h1:vtL+3mdHx/wcj3iEGz84rQa8vEqR6XM84v5Lcvfph20= modernc.org/cc/v3 v3.40.0/go.mod h1:/bTg4dnWkSXowUO6ssQKnOV0yMVxDYNIsIrzqTFDGH0= -modernc.org/cc/v4 v4.28.4 h1:Hd/4Es+MBj+/7hSdZaisNyu6bv3V0Dp2MdllyfqaH+c= -modernc.org/cc/v4 v4.28.4/go.mod h1:OnovgIhbbMXMu1aISnJ0wvVD1KnW+cAUJkIrAWh+kVI= +modernc.org/cc/v4 v4.29.0 h1:CXgwL8cvxmyzBQZzbSl/6xFtMCryb6u8IOqDci39cgc= +modernc.org/cc/v4 v4.29.0/go.mod h1:OnovgIhbbMXMu1aISnJ0wvVD1KnW+cAUJkIrAWh+kVI= modernc.org/ccgo/v3 v3.0.0-20220904174949-82d86e1b6d56/go.mod h1:YSXjPL62P2AMSxBphRHPn7IkzhVHqkvOnRKAKh+W6ZI= modernc.org/ccgo/v3 v3.16.13-0.20221017192402-261537637ce8/go.mod h1:fUB3Vn0nVPReA+7IG7yZDfjv1TMWjhQP8gCxrFAtL5g= modernc.org/ccgo/v3 v3.16.13/go.mod h1:2Quk+5YgpImhPjv2Qsob1DnZ/4som1lJTodubIcoUkY= -modernc.org/ccgo/v4 v4.34.4 h1:OVnSOWQjVKOYkFxoHYB+qQmSHK5gqMqARM+K9DpR/Ws= -modernc.org/ccgo/v4 v4.34.4/go.mod h1:qdKqE8FNIYyysougB1RX9MxCzp5oJOcQXSobANJ4TuE= +modernc.org/ccgo/v4 v4.34.6 h1:sBgfIwyN0TQ9C5hwIeuqyeAKyMWnbvj2fvpF4L11uzU= +modernc.org/ccgo/v4 v4.34.6/go.mod h1:SZ8YcN9NG7XVsQYdm6jYBvi8PQP1qi+kqB6OhjqI3Fk= modernc.org/ccorpus v1.11.6/go.mod h1:2gEUTrWqdpH2pXsmTM1ZkjeSrUWDpjMu2T6m29L/ErQ= modernc.org/fileutil v1.4.0 h1:j6ZzNTftVS054gi281TyLjHPp6CPHr2KCxEXjEbD6SM= modernc.org/fileutil v1.4.0/go.mod h1:EqdKFDxiByqxLk8ozOxObDSfcVOv/54xDs/DUHdvCUU= modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI= modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito= -modernc.org/gc/v3 v3.1.3 h1:6QAplYyVO+KdPW3pGnqmJDUxtkec8ooEWvks/hhU3lc= -modernc.org/gc/v3 v3.1.3/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY= +modernc.org/gc/v3 v3.1.4 h1:2g65LGVSmFQrXeITAw97x7hCRvZFcyE1uDP+7Vng7JI= +modernc.org/gc/v3 v3.1.4/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY= modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks= modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI= modernc.org/httpfs v1.0.6/go.mod h1:7dosgurJGp0sPaRanU53W4xZYKh14wfzX420oZADeHM= @@ -379,8 +377,6 @@ modernc.org/libc v1.20.3/go.mod h1:ZRfIaEkgrYgZDl6pa4W39HgN5G/yDW+NRmNKZBDFrk0= modernc.org/libc v1.21.4/go.mod h1:przBsL5RDOZajTVslkugzLBj1evTue36jEomFQOoYuI= modernc.org/libc v1.22.5/go.mod h1:jj+Z7dTNX8fBScMVNRAYZ/jF91K8fdT2hYMThc3YjBY= modernc.org/libc v1.24.1/go.mod h1:FmfO1RLrU3MHJfyi9eYYmZBfi/R+tqZ6+hQ3yQQUkak= -modernc.org/libc v1.73.4 h1:+ra4Ui8ngyt8HDcO1FTDPWlkAh6yOdaO2yAoh8MddQA= -modernc.org/libc v1.73.4/go.mod h1:DXZ3eO8qMCNn2SnmTNCiC71nJ9Rcq3PsnpU6Vc4rWK8= modernc.org/libc v1.74.1 h1:bdR4VTKFMC4966QSNZ05XLGI/VwzVa2kTUX51Dm0riQ= modernc.org/libc v1.74.1/go.mod h1:uH4t5bOx3G3g9Xcmj10YKlTcVISlRDwv8VoQJG9n8Os= modernc.org/mathutil v1.5.0/go.mod h1:mZW8CKdRPY1v87qxC/wUdX5O1qDzXMP5TH3wjfpga6E= @@ -399,8 +395,6 @@ modernc.org/opt v0.2.0/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns= modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w= modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE= modernc.org/sqlite v1.26.0/go.mod h1:FL3pVXie73rg3Rii6V/u5BoHlSoyeZeIgKZEgHARyCU= -modernc.org/sqlite v1.53.0 h1:20WG8N9q4ji/dEqGk4uiI0c6OPjSeLTNYGFCc3+7c1M= -modernc.org/sqlite v1.53.0/go.mod h1:xoEpOIpGrgT48H5iiyt/YXPCZPEzlfmfFwtk8Lklw8s= modernc.org/sqlite v1.54.0 h1:JCxR4qwkJvOaqAoYcgDoO25Nc+ROg6EJ2LfBVzdrgog= modernc.org/sqlite v1.54.0/go.mod h1:4ntCLuNmnH8+GNqjka1wNg7KJd5/Hi5FYp8K+XQ7GZw= modernc.org/strutil v1.1.3/go.mod h1:MEHNA7PdEnEwLvspRMtWTNnp2nnyvMfkimT1NKNAGbw= diff --git a/internal/agenttask/runtime_runner.go b/internal/agenttask/runtime_runner.go index 97cb7c23..75a7c11d 100644 --- a/internal/agenttask/runtime_runner.go +++ b/internal/agenttask/runtime_runner.go @@ -9,6 +9,7 @@ import ( "github.com/omarluq/librecode/internal/agent" "github.com/omarluq/librecode/internal/assistant" "github.com/omarluq/librecode/internal/database" + "github.com/omarluq/librecode/internal/model" ) // RuntimeRunner executes durable tasks through the shared assistant runtime. @@ -59,31 +60,62 @@ func (runner *RuntimeRunner) Run( var eventErr error - response, err := runtime.Prompt(ctx, &assistant.PromptRequest{ + metrics := new(assistant.RunMetrics) + runCtx := assistant.WithRunMetrics(ctx, metrics) + response, err := runtime.Prompt(runCtx, &assistant.PromptRequest{ OnEvent: func(event assistant.StreamEvent) { + metrics.ObserveStreamEvent(event) + if eventErr == nil { - eventErr = sink(ctx, string(event.Kind), event) + eventErr = sink(runCtx, string(event.Kind), event) } }, OnRetry: nil, OnUserEntry: nil, ParentEntryID: nil, SessionID: task.ChildSessionID, CWD: session.CWD, Text: task.Prompt, Name: "", ResumeLatest: false, HideUserPrompt: false, }) + usageJSON, usageErr := agentUsageJSON(response, metrics.Snapshot()) + if err != nil { - return Result{Text: "", UsageJSON: "{}"}, oops.In("agenttask").Code("run_prompt").Wrapf(err, "run agent prompt") + if usageErr != nil { + usageJSON = "{}" + } + + return Result{Text: "", UsageJSON: usageJSON}, + oops.In("agenttask").Code("run_prompt").Wrapf(err, "run agent prompt") } if eventErr != nil { return Result{Text: response.Text, UsageJSON: "{}"}, eventErr } - usage, err := json.Marshal(response.Usage) + if usageErr != nil { + return Result{Text: response.Text, UsageJSON: "{}"}, usageErr + } + + return Result{Text: response.Text, UsageJSON: usageJSON}, nil +} + +type agentUsage struct { + model.TokenUsage + ProviderRoundTrips int `json:"provider_round_trips,omitempty"` +} + +func agentUsageJSON(response *assistant.PromptResponse, metrics assistant.RunMetricsSnapshot) (string, error) { + usage := model.EmptyTokenUsage() + if response != nil { + usage = response.Usage + } else { + usage.InputTokens = metrics.InputTokens + usage.OutputTokens = metrics.OutputTokens + } + + encoded, err := json.Marshal(agentUsage{TokenUsage: usage, ProviderRoundTrips: metrics.ProviderRoundTrips}) if err != nil { - return Result{Text: response.Text, UsageJSON: "{}"}, oops.In("agenttask").Code("marshal_usage"). - Wrapf(err, "marshal agent usage") + return "{}", oops.In("agenttask").Code("marshal_usage").Wrapf(err, "marshal agent usage") } - return Result{Text: response.Text, UsageJSON: string(usage)}, nil + return string(encoded), nil } func (runner *RuntimeRunner) taskDefinition(task *database.AgentTaskEntity) (*agent.Definition, error) { diff --git a/internal/agenttask/runtime_runner_internal_test.go b/internal/agenttask/runtime_runner_internal_test.go index 89a90825..a8a22021 100644 --- a/internal/agenttask/runtime_runner_internal_test.go +++ b/internal/agenttask/runtime_runner_internal_test.go @@ -200,7 +200,10 @@ func runPromptScenarios( t.Context(), task, func(context.Context, string, any) error { return nil }, ) require.ErrorContains(t, err, "run agent prompt") - assert.JSONEq(t, `{}`, result.UsageJSON) + + var failedUsage agentUsage + require.NoError(t, json.Unmarshal([]byte(result.UsageJSON), &failedUsage)) + assert.Equal(t, 3, failedUsage.InputTokens) sinkErr := errors.New("persist event") result, err = newRunner(runnerCompleter{err: nil}).Run( diff --git a/internal/agenttask/service.go b/internal/agenttask/service.go index ba9211ad..59c1bc1b 100644 --- a/internal/agenttask/service.go +++ b/internal/agenttask/service.go @@ -78,6 +78,7 @@ type SubmitRequest struct { type Options struct { Tasks *database.TaskRepository AgentTasks *database.AgentTaskRepository + Workflows *database.WorkflowRepository Runner Runner Logger *slog.Logger Timeout time.Duration @@ -95,6 +96,7 @@ type Service struct { active map[string]context.CancelFunc subscribers map[string]map[uint64]chan database.TaskEventEntity agentTasks *database.AgentTaskRepository + workflows *database.WorkflowRepository queue chan string cancel context.CancelFunc done <-chan struct{} @@ -114,9 +116,13 @@ type Service struct { mu sync.Mutex } +func invalidOptions(options *Options) bool { + return options == nil || options.Tasks == nil || options.AgentTasks == nil || options.Runner == nil +} + // New creates and starts a task service. -func New(ctx context.Context, options Options) (*Service, error) { - if options.Tasks == nil || options.AgentTasks == nil || options.Runner == nil { +func New(ctx context.Context, options *Options) (*Service, error) { + if invalidOptions(options) { return nil, errors.New("agenttask: tasks, agent tasks, and runner are required") } @@ -142,7 +148,7 @@ func New(ctx context.Context, options Options) (*Service, error) { service := &Service{ runner: options.Runner, done: serviceCtx.Done(), tasks: options.Tasks, - agentTasks: options.AgentTasks, queue: make(chan string, queueCapacity), + agentTasks: options.AgentTasks, workflows: options.Workflows, queue: make(chan string, queueCapacity), cancel: cancel, active: make(map[string]context.CancelFunc), sessionSlots: make(map[string]chan struct{}), subscribers: make(map[string]map[uint64]chan database.TaskEventEntity), @@ -175,7 +181,7 @@ func New(ctx context.Context, options Options) (*Service, error) { return service, nil } -func optionDefaults(options Options) ( +func optionDefaults(options *Options) ( concurrency int, sessionConcurrency int, queueCapacity int, @@ -209,23 +215,57 @@ func (service *Service) SubmitAgentTask( ctx context.Context, request *assistant.AgentTaskRequest, ) (*database.AgentTaskEntity, error) { - return service.Submit(ctx, &SubmitRequest{ - ParentTaskID: "", - OwnerSessionID: request.OwnerSessionID, - ChildSessionID: request.ChildSessionID, - ConcurrencyKey: request.OwnerSessionID, - AgentName: request.AgentName, - Prompt: request.Prompt, - Model: request.Model, - Provider: request.Provider, - PolicyJSON: request.PolicyJSON, - Depth: request.Depth, - }) + submit := &SubmitRequest{ + ParentTaskID: request.ParentTaskID, OwnerSessionID: request.OwnerSessionID, + ChildSessionID: request.ChildSessionID, ConcurrencyKey: request.ConcurrencyKey, + AgentName: request.AgentName, Prompt: request.Prompt, Model: request.Model, + Provider: request.Provider, PolicyJSON: request.PolicyJSON, Depth: request.Depth, + } + childRequest := &database.ChildSessionRequest{ + CWD: request.ChildSessionCWD, Name: request.ChildSessionName, + ParentSessionID: request.OwnerSessionID, + } + + var created *database.AgentTaskEntity + + var err error + + if request.ParentTaskID != "" && service.workflows == nil { + return nil, oops.In("agenttask").Code("workflow_repository_missing"). + Errorf("workflow repository is required for workflow agent tasks") + } + + switch { + case request.ParentTaskID != "": + created, err = service.workflows.CreateAgentTaskWithChildSession( + ctx, request.ParentTaskID, agentTaskEntity(submit), childRequest, + request.NodeKey, request.InvocationIndex, + ) + case request.ChildSessionID == "": + created, err = service.agentTasks.CreateWithChildSession(ctx, agentTaskEntity(submit), childRequest) + default: + return service.Submit(ctx, submit) + } + + if err != nil { + return nil, oops.In("agenttask").Code("create_agent_task").Wrapf(err, "create agent task with child session") + } + + return service.enqueueCreated(ctx, created) } // Submit durably accepts a task before making it available to workers. func (service *Service) Submit(ctx context.Context, request *SubmitRequest) (*database.AgentTaskEntity, error) { - created, err := service.agentTasks.Create(ctx, &database.AgentTaskEntity{ + created, err := service.agentTasks.Create(ctx, agentTaskEntity(request)) + if err != nil { + return nil, oops.In("agenttask").Code("create_task").Wrapf(err, "create agent task") + } + + return service.enqueueCreated(ctx, created) +} + +func agentTaskEntity(request *SubmitRequest) *database.AgentTaskEntity { + return &database.AgentTaskEntity{ Task: database.TaskEntity{ ID: "", Kind: "", ParentTaskID: request.ParentTaskID, OwnerSessionID: request.OwnerSessionID, ConcurrencyKey: request.ConcurrencyKey, @@ -236,12 +276,7 @@ func (service *Service) Submit(ctx context.Context, request *SubmitRequest) (*da ChildSessionID: request.ChildSessionID, AgentName: request.AgentName, Prompt: request.Prompt, Model: request.Model, Provider: request.Provider, PolicyJSON: request.PolicyJSON, UsageJSON: "{}", Depth: request.Depth, - }) - if err != nil { - return nil, oops.In("agenttask").Code("create_task").Wrapf(err, "create agent task") } - - return service.enqueueCreated(ctx, created) } func (service *Service) enqueueCreated( diff --git a/internal/agenttask/service_internal_test.go b/internal/agenttask/service_internal_test.go index fa187a9f..254b24f8 100644 --- a/internal/agenttask/service_internal_test.go +++ b/internal/agenttask/service_internal_test.go @@ -22,7 +22,8 @@ const workerName = "worker" func emptyService() *Service { return &Service{ - runner: nil, getTaskFn: nil, renewLeaseFn: nil, active: nil, subscribers: nil, agentTasks: nil, queue: nil, + runner: nil, getTaskFn: nil, renewLeaseFn: nil, active: nil, subscribers: nil, + agentTasks: nil, workflows: nil, queue: nil, cancel: nil, done: nil, sessionSlots: nil, tasks: nil, logger: nil, leaseOwner: "", wg: sync.WaitGroup{}, nextSubscriber: 0, timeout: 0, sessionConcurrency: 0, leaseDuration: 0, leaseHeartbeatInterval: 0, leaseRenewalRetryInterval: 0, leaseRenewalAttemptTimeout: 0, @@ -134,8 +135,8 @@ func leaseRenewalService( func TestServiceInternalDefaultsAndTimeoutParsing(t *testing.T) { t.Parallel() - concurrency, sessionConcurrency, capacity, timeout := optionDefaults(Options{ - Tasks: nil, AgentTasks: nil, Runner: nil, Logger: nil, Concurrency: 7, + concurrency, sessionConcurrency, capacity, timeout := optionDefaults(&Options{ + Tasks: nil, AgentTasks: nil, Workflows: nil, Runner: nil, Logger: nil, Concurrency: 7, SessionConcurrency: 3, QueueCapacity: 9, Timeout: time.Second, }) assert.Equal(t, 7, concurrency) diff --git a/internal/agenttask/service_test.go b/internal/agenttask/service_test.go index 1f1dab44..1d859d1e 100644 --- a/internal/agenttask/service_test.go +++ b/internal/agenttask/service_test.go @@ -20,6 +20,7 @@ import ( const ( completedResult = "done" generalAgentName = "general" + testAgentPrompt = "work" ) type fakeRunner struct { @@ -72,8 +73,12 @@ func (runner *fakeRunner) unblock() { func TestNewServiceValidatesOptionsAndUsesDefaults(t *testing.T) { t.Parallel() - _, err := agenttask.New(context.Background(), agenttask.Options{ - Tasks: nil, AgentTasks: nil, Runner: nil, Logger: nil, Timeout: 0, + service, err := agenttask.New(context.Background(), nil) + require.ErrorContains(t, err, "required") + assert.Nil(t, service) + + _, err = agenttask.New(context.Background(), &agenttask.Options{ + Tasks: nil, AgentTasks: nil, Workflows: nil, Runner: nil, Logger: nil, Timeout: 0, Concurrency: 0, SessionConcurrency: 0, QueueCapacity: 0, }) require.ErrorContains(t, err, "required") @@ -82,8 +87,8 @@ func TestNewServiceValidatesOptionsAndUsesDefaults(t *testing.T) { var nilContext context.Context - _, err = agenttask.New(nilContext, agenttask.Options{ - Tasks: tasks, AgentTasks: agentTasks, Runner: new(fakeRunner), Logger: nil, Concurrency: 0, + _, err = agenttask.New(nilContext, &agenttask.Options{ + Tasks: tasks, AgentTasks: agentTasks, Workflows: nil, Runner: new(fakeRunner), Logger: nil, Concurrency: 0, SessionConcurrency: 0, QueueCapacity: 0, Timeout: 0, }) require.ErrorContains(t, err, "process context") @@ -101,8 +106,10 @@ func TestServiceAgentTaskAdaptersAndSubscriptionCancellation(t *testing.T) { }) created, err := service.SubmitAgentTask(t.Context(), &assistant.AgentTaskRequest{ - OwnerSessionID: parent.ID, ChildSessionID: child.ID, AgentName: generalAgentName, - Prompt: "work", Model: "", Provider: "", PolicyJSON: `{}`, Depth: 1, + ParentTaskID: "", OwnerSessionID: parent.ID, ChildSessionID: child.ID, + ChildSessionCWD: "", ChildSessionName: "", AgentName: generalAgentName, + Prompt: testAgentPrompt, Model: "", Provider: "", PolicyJSON: `{}`, ConcurrencyKey: parent.ID, + NodeKey: "", InvocationIndex: 0, Depth: 1, }) require.NoError(t, err) @@ -116,6 +123,27 @@ func TestServiceAgentTaskAdaptersAndSubscriptionCancellation(t *testing.T) { assert.Equal(t, database.TaskSucceeded, completed.Task.State) } +func TestServiceRejectsWorkflowSubmissionWithoutRepository(t *testing.T) { + t.Parallel() + + tasks, agentTasks, sessions := repositories(t) + owner := createSession(t, sessions, "owner", "") + service := newService(t, tasks, agentTasks, new(fakeRunner)) + + created, err := service.SubmitAgentTask(t.Context(), &assistant.AgentTaskRequest{ + ParentTaskID: "019f6830-8622-7b1a-a810-55cdac44de22", OwnerSessionID: owner.ID, ChildSessionID: "", + ChildSessionCWD: owner.CWD, ChildSessionName: "child", AgentName: generalAgentName, + Prompt: testAgentPrompt, Model: "", Provider: "", PolicyJSON: `{}`, ConcurrencyKey: owner.ID, + NodeKey: "node", InvocationIndex: 0, Depth: 1, + }) + require.ErrorContains(t, err, "workflow repository is required") + assert.Nil(t, created) + + children, listErr := sessions.ListChildSessions(t.Context(), owner.ID) + require.NoError(t, listErr) + assert.Empty(t, children) +} + func TestServiceAwaitMissingAndCanceled(t *testing.T) { t.Parallel() @@ -311,7 +339,7 @@ func TestSecondServiceDoesNotInterruptLiveTask(t *testing.T) { first := newService(t, tasks, agentTasks, firstRunner) created, err := first.Submit(t.Context(), &agenttask.SubmitRequest{ ParentTaskID: "", OwnerSessionID: parent.ID, ChildSessionID: child.ID, ConcurrencyKey: parent.ID, - AgentName: generalAgentName, Prompt: "work", Model: "", Provider: "", PolicyJSON: `{}`, Depth: 1, + AgentName: generalAgentName, Prompt: testAgentPrompt, Model: "", Provider: "", PolicyJSON: `{}`, Depth: 1, }) require.NoError(t, err) require.Equal(t, created.Task.ID, awaitStarted(t, firstRunner.started)) @@ -344,8 +372,8 @@ func TestServiceEnforcesSessionConcurrency(t *testing.T) { result: agenttask.Result{Text: completedResult, UsageJSON: `{}`}, err: nil, started: make(chan string, 3), release: make(chan struct{}), eventRelease: nil, once: sync.Once{}, } - service, err := agenttask.New(context.Background(), agenttask.Options{ - Tasks: tasks, AgentTasks: agentTasks, Runner: runner, Logger: nil, Concurrency: 2, + service, err := agenttask.New(context.Background(), &agenttask.Options{ + Tasks: tasks, AgentTasks: agentTasks, Workflows: nil, Runner: runner, Logger: nil, Concurrency: 2, SessionConcurrency: 1, QueueCapacity: 4, Timeout: time.Minute, }) require.NoError(t, err) @@ -387,8 +415,8 @@ func TestServiceRejectsWorkWhenQueueIsFull(t *testing.T) { result: agenttask.Result{Text: completedResult, UsageJSON: `{}`}, err: nil, started: make(chan string, 1), release: make(chan struct{}), eventRelease: nil, once: sync.Once{}, } - service, err := agenttask.New(context.Background(), agenttask.Options{ - Tasks: tasks, AgentTasks: agentTasks, Runner: runner, Logger: nil, Concurrency: 1, + service, err := agenttask.New(context.Background(), &agenttask.Options{ + Tasks: tasks, AgentTasks: agentTasks, Workflows: nil, Runner: runner, Logger: nil, Concurrency: 1, SessionConcurrency: 1, QueueCapacity: 1, Timeout: time.Minute, }) require.NoError(t, err) @@ -450,8 +478,8 @@ func newService( ) *agenttask.Service { t.Helper() - service, err := agenttask.New(context.Background(), agenttask.Options{ - Tasks: tasks, AgentTasks: agentTasks, Runner: runner, Logger: nil, Concurrency: 1, + service, err := agenttask.New(context.Background(), &agenttask.Options{ + Tasks: tasks, AgentTasks: agentTasks, Workflows: nil, Runner: runner, Logger: nil, Concurrency: 1, SessionConcurrency: 0, QueueCapacity: 0, Timeout: time.Minute, }) require.NoError(t, err) diff --git a/internal/assistant/agent_runtime_wrappers_internal_test.go b/internal/assistant/agent_runtime_wrappers_internal_test.go index 085b4257..68178d58 100644 --- a/internal/assistant/agent_runtime_wrappers_internal_test.go +++ b/internal/assistant/agent_runtime_wrappers_internal_test.go @@ -145,7 +145,7 @@ func TestAsyncAgentPromptAndToolRegistries(t *testing.T) { assert.Contains(t, prompt, "Use agent_start") assert.Contains(t, prompt, "agent_cancel") - registry, err := runtime.promptToolRegistry(cwd, "owner") + registry, err := runtime.promptToolRegistry(t.Context(), cwd, "owner") require.NoError(t, err) names := make([]tool.Name, 0, len(registry.Definitions())) @@ -160,7 +160,7 @@ func TestAsyncAgentPromptAndToolRegistries(t *testing.T) { } runtime.agentTasks = nil - registry, err = runtime.promptToolRegistry(cwd, "owner") + registry, err = runtime.promptToolRegistry(t.Context(), cwd, "owner") require.NoError(t, err) for _, definition := range registry.Definitions() { @@ -168,6 +168,35 @@ func TestAsyncAgentPromptAndToolRegistries(t *testing.T) { } } +func TestPromptToolRegistryHonorsToolStrategy(t *testing.T) { + t.Parallel() + + tests := []struct { + strategy ToolStrategy + name string + wantExecute bool + }{ + {name: "default hybrid", strategy: "", wantExecute: true}, + {name: "hybrid", strategy: ToolStrategyHybrid, wantExecute: true}, + {name: "direct", strategy: ToolStrategyDirect, wantExecute: false}, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + runtime := new(Runtime) + runtime.profile = topLevelExecutionProfile() + registry, err := runtime.promptToolRegistry( + WithToolStrategy(t.Context(), testCase.strategy), t.TempDir(), "owner", + ) + require.NoError(t, err) + + assert.Equal(t, testCase.wantExecute, registry.Has(executeToolName)) + }) + } +} + func TestAgentExecutionProfileRegistryAndPrompt(t *testing.T) { t.Parallel() cwd := t.TempDir() @@ -182,7 +211,7 @@ func TestAgentExecutionProfileRegistryAndPrompt(t *testing.T) { assert.Contains(t, prompt, "Use only the available tools (read, write)") assert.NotContains(t, prompt, "Use agent_start") - registry, err := runtime.promptToolRegistry(cwd, "ignored") + registry, err := runtime.promptToolRegistry(t.Context(), cwd, "ignored") require.NoError(t, err) _, err = registry.Execute(t.Context(), string(tool.NameWrite), agentArguments(t, `{"path":"x","content":"y"}`)) require.ErrorContains(t, err, "mutating tool is denied") @@ -219,7 +248,7 @@ func TestToolRegistryProviderAndCollisionBranches(t *testing.T) { assert.Equal(t, "ok", result.Text()) provider.tools = []extension.Tool{{ - Name: "read", Description: "collision", Extension: "", InputSchema: tool.EmptySchema(), + Name: jsonReadToolName, Description: "collision", Extension: "", InputSchema: tool.EmptySchema(), }} _, err = newToolRegistry(t.TempDir(), provider) require.ErrorContains(t, err, "register extension tools") @@ -232,7 +261,7 @@ func TestToolRegistryProviderAndCollisionBranches(t *testing.T) { runtime.extensions = testToolProvider{tools: []extension.Tool{{ Name: string(agentStartToolName), Description: "collision", Extension: "", InputSchema: tool.EmptySchema(), }}} - _, err = runtime.promptToolRegistry(t.TempDir(), "owner") + _, err = runtime.promptToolRegistry(t.Context(), t.TempDir(), "owner") require.ErrorContains(t, err, "register agent tool") } @@ -245,7 +274,7 @@ func TestPromptRegistryHonorsDisabledExtensions(t *testing.T) { }}} runtime.profile = topLevelExecutionProfile() runtime.profile.EnableExtensions = false - registry, err := runtime.promptToolRegistry(t.TempDir(), "owner") + registry, err := runtime.promptToolRegistry(t.Context(), t.TempDir(), "owner") require.NoError(t, err) for _, definition := range registry.Definitions() { diff --git a/internal/assistant/agent_submitter.go b/internal/assistant/agent_submitter.go new file mode 100644 index 00000000..6d6135de --- /dev/null +++ b/internal/assistant/agent_submitter.go @@ -0,0 +1,115 @@ +package assistant + +import ( + "context" + "encoding/json" + "strings" + + "github.com/samber/oops" + + "github.com/omarluq/librecode/internal/agent" + "github.com/omarluq/librecode/internal/database" +) + +const defaultWorkflowAgentName = "general" + +// AgentSubmitRequest describes one high-level durable agent launch. +type AgentSubmitRequest struct { + ParentTaskID string + OwnerSessionID string + CWD string + AgentName string + Prompt string + Model string + Provider string + ConcurrencyKey string + NodeKey string + InvocationIndex int + Depth int +} + +// AgentSubmitter creates child sessions and snapshots agent policy before +// submitting durable work. +type AgentSubmitter struct { + controller AgentTaskController + catalog *agent.Catalog +} + +// NewAgentSubmitter creates the shared high-level agent submission boundary. +func NewAgentSubmitter( + controller AgentTaskController, + catalog *agent.Catalog, +) (*AgentSubmitter, error) { + if controller == nil || catalog == nil { + return nil, oops.In("assistant").Code("invalid_agent_submitter_dependencies"). + Errorf("agent task controller and catalog are required") + } + + return &AgentSubmitter{controller: controller, catalog: catalog}, nil +} + +// SubmitAgent resolves policy, creates a child session, and durably submits work. +func (submitter *AgentSubmitter) SubmitAgent( + ctx context.Context, + request *AgentSubmitRequest, +) (*database.AgentTaskEntity, error) { + if request == nil { + return nil, oops.In("assistant").Code("nil_agent_submit_request").Errorf("agent submit request is nil") + } + + definition, err := submitter.resolveDefinition(request) + if err != nil { + return nil, err + } + + policy, err := json.Marshal(&definition) + if err != nil { + return nil, oops.In("assistant").Code("encode_agent_profile").Wrapf(err, "encode agent profile") + } + + concurrencyKey := request.ConcurrencyKey + if concurrencyKey == "" { + concurrencyKey = request.OwnerSessionID + } + + depth := request.Depth + if depth <= 0 { + depth = 1 + } + + task, err := submitter.controller.SubmitAgentTask(ctx, &AgentTaskRequest{ + ParentTaskID: request.ParentTaskID, OwnerSessionID: request.OwnerSessionID, + ChildSessionID: "", ChildSessionCWD: request.CWD, + ChildSessionName: childSessionName(definition.Name, request.Prompt), AgentName: definition.Name, + Prompt: request.Prompt, Model: definition.Model.Model, Provider: definition.Model.Provider, + PolicyJSON: string(policy), ConcurrencyKey: concurrencyKey, + NodeKey: request.NodeKey, InvocationIndex: request.InvocationIndex, Depth: depth, + }) + if err != nil { + return task, oops.In("assistant").Code("submit_agent_task").Wrapf(err, "submit agent task") + } + + return task, nil +} + +func (submitter *AgentSubmitter) resolveDefinition(request *AgentSubmitRequest) (agent.Definition, error) { + agentName := strings.ToLower(strings.TrimSpace(request.AgentName)) + if agentName == "" { + agentName = defaultWorkflowAgentName + } + + definition, found := submitter.catalog.Get(agentName) + if !found { + return agent.Definition{}, oops.In("assistant").Code("unknown_agent").Errorf("unknown agent %q", agentName) + } + + if request.Model != "" { + definition.Model.Model = request.Model + } + + if request.Provider != "" { + definition.Model.Provider = request.Provider + } + + return definition, nil +} diff --git a/internal/assistant/agent_submitter_internal_test.go b/internal/assistant/agent_submitter_internal_test.go new file mode 100644 index 00000000..4108b3f7 --- /dev/null +++ b/internal/assistant/agent_submitter_internal_test.go @@ -0,0 +1,65 @@ +package assistant + +import ( + "encoding/json" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/omarluq/librecode/internal/agent" + "github.com/omarluq/librecode/internal/database" +) + +const agentSubmitterOwnerID = "owner" + +func TestAgentSubmitterBuildsAtomicChildSubmissionAndSnapshotsPolicy(t *testing.T) { + t.Parallel() + + controller := newAgentControllerStub(agentToolTask("task-1", "owner", database.TaskQueued), nil, true) + cwd := t.TempDir() + submitter, err := NewAgentSubmitter(controller, isolatedAgentCatalog(t)) + require.NoError(t, err) + + task, err := submitter.SubmitAgent(t.Context(), &AgentSubmitRequest{ + ParentTaskID: "", OwnerSessionID: agentSubmitterOwnerID, CWD: cwd, AgentName: "", Prompt: "inspect code", + Model: "model-override", Provider: "provider-override", ConcurrencyKey: agentSubmitterOwnerID, + NodeKey: "", InvocationIndex: 0, Depth: 2, + }) + require.NoError(t, err) + require.NotNil(t, task) + require.NotNil(t, controller.lastSubmit) + assert.Equal(t, defaultWorkflowAgentName, controller.lastSubmit.AgentName) + assert.Equal(t, "model-override", controller.lastSubmit.Model) + assert.Equal(t, "provider-override", controller.lastSubmit.Provider) + assert.Equal(t, 2, controller.lastSubmit.Depth) + assert.Empty(t, controller.lastSubmit.ChildSessionID) + assert.Equal(t, cwd, controller.lastSubmit.ChildSessionCWD) + assert.NotEmpty(t, controller.lastSubmit.ChildSessionName) + + var definition agent.Definition + require.NoError(t, json.Unmarshal([]byte(controller.lastSubmit.PolicyJSON), &definition)) + assert.Equal(t, defaultWorkflowAgentName, definition.Name) + assert.Equal(t, "model-override", definition.Model.Model) + assert.Equal(t, "provider-override", definition.Model.Provider) +} + +func TestAgentSubmitterReturnsRejectedAtomicSubmission(t *testing.T) { + t.Parallel() + + controller := newAgentControllerStub(nil, nil, false) + controller.submitErr = errors.New("rejected") + submitter, err := NewAgentSubmitter(controller, isolatedAgentCatalog(t)) + require.NoError(t, err) + + _, err = submitter.SubmitAgent(t.Context(), &AgentSubmitRequest{ + ParentTaskID: "", OwnerSessionID: agentSubmitterOwnerID, CWD: t.TempDir(), + AgentName: defaultWorkflowAgentName, Prompt: "inspect", Model: "", Provider: "", + ConcurrencyKey: agentSubmitterOwnerID, NodeKey: "", InvocationIndex: 0, Depth: 1, + }) + require.Error(t, err) + require.NotNil(t, controller.lastSubmit) + assert.Empty(t, controller.lastSubmit.ChildSessionID) + assert.NotEmpty(t, controller.lastSubmit.ChildSessionName) +} diff --git a/internal/assistant/agent_tool.go b/internal/assistant/agent_tool.go index 3d0d1fb6..3df4f65b 100644 --- a/internal/assistant/agent_tool.go +++ b/internal/assistant/agent_tool.go @@ -2,7 +2,6 @@ package assistant import ( "context" - "encoding/json" "errors" "fmt" "strings" @@ -29,14 +28,20 @@ const ( // AgentTaskRequest describes one asynchronous agent submission. type AgentTaskRequest struct { - OwnerSessionID string - ChildSessionID string - AgentName string - Prompt string - Model string - Provider string - PolicyJSON string - Depth int + ParentTaskID string + OwnerSessionID string + ChildSessionID string + ChildSessionCWD string + ChildSessionName string + AgentName string + Prompt string + Model string + Provider string + PolicyJSON string + ConcurrencyKey string + NodeKey string + InvocationIndex int + Depth int } // AgentTaskController is the runtime-facing boundary for durable agent work. @@ -230,36 +235,23 @@ func (executor *agentToolExecutor) submit( definition *agent.Definition, prompt string, ) (*database.AgentTaskEntity, error) { - policy, err := json.Marshal(definition) + submitter, err := NewAgentSubmitter(executor.controller, executor.catalog) if err != nil { - return nil, oops.In("assistant").Code("encode_agent_profile").Wrapf(err, "encode agent profile") - } - - child, err := executor.sessions.CreateSession( - ctx, executor.cwd, childSessionName(definition.Name, prompt), executor.parentSessionID, - ) - if err != nil { - return nil, oops.In("assistant").Code("create_agent_session").Wrapf(err, "create child agent session") - } - - task, err := executor.controller.SubmitAgentTask(ctx, &AgentTaskRequest{ - OwnerSessionID: executor.parentSessionID, ChildSessionID: child.ID, AgentName: definition.Name, - Prompt: prompt, Model: definition.Model.Model, Provider: definition.Model.Provider, - PolicyJSON: string(policy), Depth: 1, + return nil, err + } + + return submitter.SubmitAgent(ctx, &AgentSubmitRequest{ + ParentTaskID: "", OwnerSessionID: executor.parentSessionID, + CWD: executor.cwd, + AgentName: definition.Name, + Prompt: prompt, + Model: definition.Model.Model, + Provider: definition.Model.Provider, + ConcurrencyKey: executor.parentSessionID, + NodeKey: "", + InvocationIndex: 0, + Depth: 1, }) - if err == nil { - return task, nil - } - - if task != nil { - return task, oops.In("assistant").Code("submit_agent_task").Wrapf(err, "submit agent task") - } - - if deleteErr := executor.sessions.DeleteSession(context.WithoutCancel(ctx), child.ID); deleteErr != nil { - return nil, errors.Join(err, deleteErr) - } - - return nil, oops.In("assistant").Code("submit_agent_task").Wrapf(err, "submit agent task") } func (executor *agentToolExecutor) status(ctx context.Context, input tool.Arguments) (tool.Result, error) { diff --git a/internal/assistant/agent_tool_internal_test.go b/internal/assistant/agent_tool_internal_test.go index 8c2ae65c..ef18d7a8 100644 --- a/internal/assistant/agent_tool_internal_test.go +++ b/internal/assistant/agent_tool_internal_test.go @@ -119,10 +119,9 @@ func TestAgentStartCreatesChildAndSubmitsSnapshot(t *testing.T) { assert.Equal(t, "inspect code", stub.lastSubmit.Prompt) assert.Equal(t, parent.ID, stub.lastSubmit.OwnerSessionID) assert.Contains(t, stub.lastSubmit.PolicyJSON, `"name":"general"`) - child, found, getErr := sessions.GetSession(t.Context(), stub.lastSubmit.ChildSessionID) - require.NoError(t, getErr) - require.True(t, found) - assert.Equal(t, parent.ID, child.ParentSession) + assert.Empty(t, stub.lastSubmit.ChildSessionID) + assert.Equal(t, parent.CWD, stub.lastSubmit.ChildSessionCWD) + assert.NotEmpty(t, stub.lastSubmit.ChildSessionName) } func TestAgentStartValidationAndSubmissionCleanup(t *testing.T) { @@ -164,9 +163,8 @@ func TestAgentStartValidationAndSubmissionCleanup(t *testing.T) { executor := newAgentToolExecutor(stub, sessions, catalog, agentStartToolName, parent.ID, parent.CWD) _, err = executor.Execute(t.Context(), agentArguments(t, `{"agent":"general","prompt":"work"}`)) require.ErrorContains(t, err, "submit agent task") - _, found, getErr := sessions.GetSession(t.Context(), stub.lastSubmit.ChildSessionID) - require.NoError(t, getErr) - assert.False(t, found) + assert.Empty(t, stub.lastSubmit.ChildSessionID) + assert.NotEmpty(t, stub.lastSubmit.ChildSessionName) } func TestAgentTaskOperationsAreOwnerScoped(t *testing.T) { @@ -319,28 +317,11 @@ func TestAgentSubmitFailureBranches(t *testing.T) { assert.Empty(t, result.Text()) assert.Equal(t, "accepted", result.Details["task_id"]) - closedSessions, databaseConnection := agentToolSessionsWithDB(t) - parent, err = closedSessions.CreateSession(t.Context(), t.TempDir(), "parent", "") - require.NoError(t, err) - stub = new(agentControllerStub) stub.submitErr = errors.New("submit failed") - stub.submitFunc = func() (*database.AgentTaskEntity, error) { - require.NoError(t, databaseConnection.Close()) - - return nil, stub.submitErr - } - executor = newAgentToolExecutor(stub, closedSessions, catalog, agentStartToolName, parent.ID, parent.CWD) + executor = newAgentToolExecutor(stub, sessions, catalog, agentStartToolName, parent.ID, parent.CWD) _, err = executor.Execute(t.Context(), agentArguments(t, `{"agent":"general","prompt":"work"}`)) require.ErrorContains(t, err, "submit failed") - require.ErrorContains(t, err, "database is closed") - - closedSessions, databaseConnection = agentToolSessionsWithDB(t) - require.NoError(t, databaseConnection.Close()) - - executor.sessions = closedSessions - _, err = executor.Execute(t.Context(), agentArguments(t, `{"agent":"general","prompt":"work"}`)) - require.ErrorContains(t, err, "create child agent session") } func TestChildSessionNameNormalizesAndTruncates(t *testing.T) { diff --git a/internal/assistant/catalog_testmain_test.go b/internal/assistant/catalog_testmain_test.go index c9aeae56..c211781f 100644 --- a/internal/assistant/catalog_testmain_test.go +++ b/internal/assistant/catalog_testmain_test.go @@ -1,11 +1,26 @@ package assistant_test import ( + "fmt" "os" "testing" + + "github.com/omarluq/librecode/internal/executeworker" ) func TestMain(m *testing.M) { + if len(os.Args) == 2 && os.Args[1] == "__execute-worker" { + if err := executeworker.Serve(os.Stdin, os.Stdout); err != nil { + if _, writeErr := fmt.Fprintln(os.Stderr, err); writeErr != nil { + os.Exit(1) + } + + os.Exit(1) + } + + os.Exit(0) + } + home, err := os.MkdirTemp("", "librecode-assistant-test-home-*") if err != nil { panic(err) diff --git a/internal/assistant/client.go b/internal/assistant/client.go index 941d8100..fb82698d 100644 --- a/internal/assistant/client.go +++ b/internal/assistant/client.go @@ -62,17 +62,22 @@ type CompletionResult struct { type ToolCallEvent struct { ArgumentsJSON string `json:"arguments_json"` ID string `json:"id"` + ParentCallID string `json:"parent_call_id,omitempty"` Name string `json:"name"` Arguments tool.Arguments `json:"arguments,omitzero"` + Sequence int `json:"sequence,omitempty"` } // ToolEvent captures one tool call for persistence and TUI rendering. type ToolEvent struct { + CallID string `json:"call_id,omitempty"` + ParentCallID string `json:"parent_call_id,omitempty"` Name string `json:"name"` ArgumentsJSON string `json:"arguments_json"` DetailsJSON string `json:"details_json,omitempty"` Result string `json:"result"` Error string `json:"error,omitempty"` + Sequence int `json:"sequence,omitempty"` IsError bool `json:"is_error,omitempty"` } diff --git a/internal/assistant/client_adapter.go b/internal/assistant/client_adapter.go index bea1b0f5..8100ae0a 100644 --- a/internal/assistant/client_adapter.go +++ b/internal/assistant/client_adapter.go @@ -2,6 +2,7 @@ package assistant import ( "context" + "math" "strings" "github.com/samber/lo" @@ -14,6 +15,11 @@ import ( "github.com/omarluq/librecode/internal/tool" ) +const ( + toolParentCallIDMetadataKey = "parent_call_id" + toolSequenceMetadataKey = "sequence" +) + func providerRequestFromCompletionRequest(request *CompletionRequest) *provider.CompletionRequest { if request == nil { return emptyProviderRequest() @@ -188,11 +194,14 @@ func toolEventsFromLLMParts(parts []llm.Part) []ToolEvent { func emptyToolEvent() ToolEvent { return ToolEvent{ + CallID: "", + ParentCallID: "", Name: "", ArgumentsJSON: "", DetailsJSON: "", Result: "", Error: "", + Sequence: 0, IsError: false, } } @@ -203,11 +212,14 @@ func toolEventFromLLMToolResult(result *llm.ToolResult) ToolEvent { } return ToolEvent{ + CallID: result.ToolCallID, + ParentCallID: stringFromOptions(result.Metadata, toolParentCallIDMetadataKey), Name: result.Name, ArgumentsJSON: result.ArgumentsJSON, DetailsJSON: stringFromOptions(result.Metadata, "details_json"), Result: textFromLLMParts(result.Content), Error: result.Error, + Sequence: sequenceFromOptions(result.Metadata), IsError: result.IsError, } } @@ -296,14 +308,23 @@ func toolCallEventPointerFromLLMPart(part *llm.Part) *ToolCallEvent { func toolCallEventFromLLMToolCall(call *llm.ToolCall) ToolCallEvent { if call == nil { - return ToolCallEvent{Arguments: tool.EmptyArguments(), ID: "", Name: "", ArgumentsJSON: ""} + return ToolCallEvent{ + ArgumentsJSON: "", + ID: "", + ParentCallID: "", + Name: "", + Arguments: tool.EmptyArguments(), + Sequence: 0, + } } return ToolCallEvent{ Arguments: call.Arguments, ID: call.ID, + ParentCallID: stringFromOptions(call.Metadata, toolParentCallIDMetadataKey), Name: call.Name, ArgumentsJSON: call.ArgumentsJSON, + Sequence: sequenceFromOptions(call.Metadata), } } @@ -319,7 +340,7 @@ func llmToolCallFromToolCallEvent(event *ToolCallEvent, fallbackName string) llm } return llm.ToolCall{ - Metadata: nil, + Metadata: toolIdentityMetadata(event.ParentCallID, event.Sequence), Arguments: event.Arguments, ID: event.ID, Name: event.Name, @@ -479,6 +500,45 @@ func stringFromOptions(options map[string]any, key string) string { return value } +func sequenceFromOptions(options map[string]any) int { + value, ok := optionInt64(options[toolSequenceMetadataKey]) + if !ok || value < 0 || uint64(value) > uint64(^uint(0)>>1) { + return 0 + } + + return int(value) +} + +func optionInt64(value any) (int64, bool) { + switch typed := value.(type) { + case int: + return int64(typed), true + case int32: + return int64(typed), true + case int64: + return typed, true + case float64: + if typed < math.MinInt64 || typed >= float64(math.MaxInt64) || math.Trunc(typed) != typed { + return 0, false + } + + return int64(typed), true + default: + return 0, false + } +} + +func toolIdentityMetadata(parentCallID string, sequence int) map[string]any { + if parentCallID == "" && sequence == 0 { + return nil + } + + return map[string]any{ + toolParentCallIDMetadataKey: parentCallID, + toolSequenceMetadataKey: sequence, + } +} + func hookAttempt(input *llm.HookInput) int { if input == nil { return 0 diff --git a/internal/assistant/client_adapter_internal_test.go b/internal/assistant/client_adapter_internal_test.go index a977256f..e5b65cbb 100644 --- a/internal/assistant/client_adapter_internal_test.go +++ b/internal/assistant/client_adapter_internal_test.go @@ -3,6 +3,8 @@ package assistant import ( "context" "errors" + "math" + "strconv" "testing" "github.com/stretchr/testify/assert" @@ -21,6 +23,8 @@ const ( adapterHeaderKey = "x-test" adapterBaseURL = "https://example.test" adapterHeaderValue = "enabled" + adapterToolCallID = "call-1" + adapterParentCallID = "outer" adapterThought = " thought " adapterBlankThought = " " adapterReadPath = "README.md" @@ -243,38 +247,55 @@ func TestToolExecutorAdapterConvertsCallsEventsAndErrors(t *testing.T) { require.Len(t, calls, 1) observedCall = calls[0] - onEvent(streamEvent(StreamEventToolResult, "", adapterToolEvent(), nil)) + event := ToolEvent{ + DetailsJSON: "", + Error: "", + IsError: false, - return []ToolEvent{{ + CallID: adapterToolCallID, + ParentCallID: adapterParentCallID, Name: jsonReadToolName, ArgumentsJSON: adapterReadArgs, - DetailsJSON: "", Result: "ok", - Error: "", - IsError: false, - }}, nil + Sequence: 2, + } + onEvent(streamEvent(StreamEventToolResult, "", &event, nil)) + + return []ToolEvent{event}, nil }, ) results, err := executor( context.Background(), []llm.ToolCall{{ - Metadata: map[string]any{"m": true}, + Metadata: map[string]any{ + "m": true, + "parent_call_id": adapterParentCallID, + toolSequenceMetadataKey: float64(2), + }, Arguments: testutil.ToolArguments(map[string]any{jsonPathKey: adapterReadPath}), - ID: "call-1", + ID: adapterToolCallID, Name: jsonReadToolName, ArgumentsJSON: adapterReadArgs, }}, func(chunk *llm.StreamChunk) { observedEvent = chunk }, ) require.NoError(t, err) - assert.Equal(t, "call-1", observedCall.ID) + assert.Equal(t, adapterToolCallID, observedCall.ID) + assert.Equal(t, adapterParentCallID, observedCall.Metadata["parent_call_id"]) assert.Equal(t, adapterReadPath, testutil.ToolArgumentFields(observedCall.Arguments)[jsonPathKey]) require.NotNil(t, observedEvent) require.NotNil(t, observedEvent.Part) assert.Equal(t, llm.PartToolResult, observedEvent.Part.Type) + require.NotNil(t, observedEvent.Part.ToolResult) + assert.Equal(t, adapterToolCallID, observedEvent.Part.ToolResult.ToolCallID) + assert.Equal(t, adapterParentCallID, observedEvent.Part.ToolResult.Metadata["parent_call_id"]) + assert.Equal(t, 2, observedEvent.Part.ToolResult.Metadata[toolSequenceMetadataKey]) require.Len(t, results, 1) assert.Equal(t, expectedReadToolName, results[0].Name) + assert.Equal(t, adapterToolCallID, results[0].ToolCallID) + assert.Equal(t, adapterParentCallID, results[0].Metadata["parent_call_id"]) + assert.Equal(t, 2, results[0].Metadata[toolSequenceMetadataKey]) assert.Nil(t, llmToolExecutor(nil)) assert.Nil(t, assistantStreamEventHandler(nil)) @@ -282,6 +303,50 @@ func TestToolExecutorAdapterConvertsCallsEventsAndErrors(t *testing.T) { assert.Nil(t, assistantToolCallsFromLLM(nil)) } +func TestIntFromOptionsStrictSequenceConversion(t *testing.T) { + t.Parallel() + + maxInt := int(^uint(0) >> 1) + + tests := []struct { + value any + name string + want int + }{ + {name: "int", value: maxInt, want: maxInt}, + {name: "int32", value: int32(math.MaxInt32), want: int(math.MaxInt32)}, + {name: "int64 in platform range", value: int64(maxInt), want: maxInt}, + {name: "integral float64", value: float64(42), want: 42}, + {name: "fractional float64", value: 1.5, want: 0}, + {name: "negative int", value: -1, want: 0}, + {name: "negative int32", value: int32(-1), want: 0}, + {name: "negative int64", value: int64(-1), want: 0}, + {name: "negative float64", value: float64(-1), want: 0}, + {name: "platform overflow float64", value: math.Ldexp(1, strconv.IntSize-1), want: 0}, + {name: "positive infinity", value: math.Inf(1), want: 0}, + {name: "not a number", value: math.NaN(), want: 0}, + {name: "unsupported type", value: "1", want: 0}, + } + if strconv.IntSize == 32 { + tests = append(tests, struct { + value any + name string + want int + }{name: "int64 outside platform range", value: int64(math.MaxInt32) + 1, want: 0}) + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + options := map[string]any{toolSequenceMetadataKey: test.value} + assert.Equal(t, test.want, sequenceFromOptions(options)) + }) + } + + assert.Zero(t, sequenceFromOptions(nil)) +} + func TestAdapterNilAndFallbackHelpers(t *testing.T) { t.Parallel() @@ -364,6 +429,10 @@ func streamEvent(kind StreamEventKind, text string, event *ToolEvent, usage *mod func adapterToolEvent() *ToolEvent { return &ToolEvent{ + CallID: "", + ParentCallID: "", + Sequence: 0, + Name: jsonReadToolName, ArgumentsJSON: "", DetailsJSON: "", diff --git a/internal/assistant/context_auto_compaction.go b/internal/assistant/context_auto_compaction.go index 9a4cb0f9..3cdc5e77 100644 --- a/internal/assistant/context_auto_compaction.go +++ b/internal/assistant/context_auto_compaction.go @@ -40,7 +40,7 @@ func (runtime *Runtime) buildCompletionRequest( return nil, oops.In("assistant").Code("context_build_model").Wrapf(err, "context: build model context") } - registry, err := runtime.promptToolRegistry(cwd, sessionID) + registry, err := runtime.promptToolRegistry(ctx, cwd, sessionID) if err != nil { return nil, oops.In("assistant").Code("context_tool_registry").Wrapf(err, "context: create tool registry") } diff --git a/internal/assistant/execute_tool.go b/internal/assistant/execute_tool.go new file mode 100644 index 00000000..9c5fafef --- /dev/null +++ b/internal/assistant/execute_tool.go @@ -0,0 +1,256 @@ +package assistant + +import ( + "context" + "encoding/json" + "strings" + + "github.com/samber/oops" + + "github.com/omarluq/librecode/internal/executeworker" + "github.com/omarluq/librecode/internal/mvmhost" + "github.com/omarluq/librecode/internal/tool" +) + +const executeToolName tool.Name = "execute" + +const ( + defaultExecuteResultLimit = 1 << 20 + executeNameKey = "name" + executeResultValueKey = "result_value" + executeCallMethod = "call" +) + +type nestedToolInvoker func(context.Context, string, tool.Arguments, string) (tool.Result, ToolEvent) + +type executeToolExecutor struct { + registry *tool.Registry + invoke nestedToolInvoker +} + +type executeToolInput struct { + Source string `json:"source"` +} + +type executeToolCallResult = executeworker.ToolCallResult + +func newExecuteTool(registry *tool.Registry, invoke nestedToolInvoker) *executeToolExecutor { + return &executeToolExecutor{registry: registry, invoke: invoke} +} + +func (executor *executeToolExecutor) Definition() tool.Definition { + return tool.Definition{ + Schema: mustToolSchema( + `{"type":"object","additionalProperties":false,"properties":` + + `{"source":{"type":"string","description":"Go source to evaluate with the tools package."}},` + + `"required":["source"]}`, + ), + Name: executeToolName, + Label: "Execute Go", + Description: "Evaluate Go source that can search, describe, and call the tools available for this prompt.", + PromptSnippet: "Use execute for compact multi-tool programs", + PromptGuidelines: []string{ + `Import "tools" to use tools.Search(query), tools.Describe(name), and tools.Call(name, input).`, + "The execute tool cannot search for, describe, or call itself.", + }, + ReadOnly: false, + } +} + +func (executor *executeToolExecutor) Execute(ctx context.Context, input tool.Arguments) (tool.Result, error) { + if executor.registry == nil { + return tool.Result{}, oops.In("assistant").Code("execute_registry_missing"). + Errorf("execute tool registry is not configured") + } + + var args executeToolInput + if err := input.Decode(&args); err != nil { + return tool.Result{}, oops.In("assistant").Code("execute_input").Wrapf(err, "decode execute input") + } + + client := executeworker.Client{Executable: "", Handler: executor.handleWorkerMessage} + + result, err := client.Eval(ctx, args.Source) + if err != nil { + wrapped := oops.In("assistant").Code("execute_source").Wrapf(err, "execute MVM source") + + return tool.TextResult("", executeResultDetails(result)), wrapped + } + + if nested, ok := result.Value.(executeworker.ToolCallResult); ok && !nested.IsError { + toolResult := tool.Result{Details: nested.Details, Content: nested.Content} + if toolResult.Details == nil { + toolResult.Details = map[string]any{} + } + + toolResult.Details["execute_stdout"] = result.Stdout + toolResult.Details["execute_stderr"] = result.Stderr + + return toolResult, nil + } + + text, err := executeResultText(result) + if err != nil { + return tool.Result{}, err + } + + return tool.TextResult(text, executeResultDetails(result)), nil +} + +func (executor *executeToolExecutor) handleWorkerMessage( + ctx context.Context, + message *executeworker.Message, +) (any, error) { + switch message.Method { + case "search": + return executor.search(message.Query), nil + case "describe": + return executor.describe(message.Name), nil + case executeCallMethod: + if _, err := tool.ArgumentsFromRaw(message.Input); err != nil { + return nil, oops.In("assistant").Code("execute_rpc_input").Wrapf(err, "decode nested tool input") + } + + return executor.call(ctx, message.Name, message.Input), nil + default: + return nil, oops.In("assistant").Code("execute_rpc_method").Errorf( + "unknown execute worker RPC method %q", + message.Method, + ) + } +} + +func (executor *executeToolExecutor) search(query string) []map[string]any { + query = strings.ToLower(strings.TrimSpace(query)) + matches := make([]map[string]any, 0) + + for _, definition := range executor.registry.Definitions() { + if definition.Name == executeToolName || definition.Name == workflowToolName { + continue + } + + haystack := strings.ToLower(strings.Join([]string{ + string(definition.Name), definition.Label, definition.Description, definition.PromptSnippet, + }, " ")) + if query == "" || strings.Contains(haystack, query) { + matches = append(matches, executeDefinitionMap(&definition)) + } + } + + return matches +} + +func (executor *executeToolExecutor) describe(name string) map[string]any { + for _, definition := range executor.registry.Definitions() { + if definition.Name == executeToolName || definition.Name == workflowToolName { + continue + } + + if string(definition.Name) == name { + return executeDefinitionMap(&definition) + } + } + + return nil +} + +func (executor *executeToolExecutor) call( + ctx context.Context, + name string, + encoded json.RawMessage, +) executeToolCallResult { + if tool.Name(name) == executeToolName { + return executeToolCallResult{ + Details: map[string]any{}, Content: nil, Error: "execute cannot call itself", IsError: true, + } + } + + if tool.Name(name) == workflowToolName { + return executeToolCallResult{ + Details: map[string]any{}, Content: nil, Error: "execute cannot call workflow", IsError: true, + } + } + + arguments, err := tool.ArgumentsFromRaw(encoded) + if err != nil { + return executeToolCallResult{Details: map[string]any{}, Content: nil, Error: err.Error(), IsError: true} + } + + var result tool.Result + + if executor.invoke != nil { + var event ToolEvent + + result, event = executor.invoke(ctx, name, arguments, string(encoded)) + + if event.IsError { + return executeToolCallResult{ + Details: result.Details, Content: result.Content, Error: event.Error, IsError: true, + } + } + } else { + result, err = executor.registry.Execute(ctx, name, arguments) + } + + outcome := executeToolCallResult{Details: result.Details, Content: result.Content, Error: "", IsError: false} + if outcome.Details == nil { + outcome.Details = map[string]any{} + } + + if err != nil { + outcome.Error = err.Error() + outcome.IsError = true + } + + return outcome +} + +func executeDefinitionMap(definition *tool.Definition) map[string]any { + var schema any + if raw := definition.Schema.RawMessage(); len(raw) > 0 { + if err := json.Unmarshal(raw, &schema); err != nil { + schema = string(raw) + } + } + + return map[string]any{ + executeNameKey: string(definition.Name), + "label": definition.Label, + "description": definition.Description, + "prompt_snippet": definition.PromptSnippet, + "prompt_guidelines": definition.PromptGuidelines, + "read_only": definition.ReadOnly, + "schema": schema, + } +} + +func executeResultText(result mvmhost.Result) (string, error) { + if result.Value == nil { + if result.Stdout != "" { + return result.Stdout, nil + } + + return "null", nil + } + + encoded, err := json.Marshal(result.Value) + if err != nil { + return "", oops.In("assistant").Code("execute_result_encode").Wrapf(err, "encode execute result") + } + + if len(encoded) > defaultExecuteResultLimit { + return "", oops.In("assistant").Code("execute_result_limit").Errorf( + "encoded execute result is %d bytes; limit is %d", + len(encoded), + defaultExecuteResultLimit, + ) + } + + return string(encoded), nil +} + +func executeResultDetails(result mvmhost.Result) map[string]any { + return map[string]any{executeResultValueKey: result.Value, "stdout": result.Stdout, "stderr": result.Stderr} +} + +var _ tool.Executor = (*executeToolExecutor)(nil) diff --git a/internal/assistant/execute_tool_internal_test.go b/internal/assistant/execute_tool_internal_test.go new file mode 100644 index 00000000..64181016 --- /dev/null +++ b/internal/assistant/execute_tool_internal_test.go @@ -0,0 +1,499 @@ +package assistant + +import ( + "context" + "encoding/json" + "strings" + "sync" + "testing" + "time" + + "github.com/samber/oops" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/omarluq/librecode/internal/executeworker" + "github.com/omarluq/librecode/internal/extension" + "github.com/omarluq/librecode/internal/mvmhost" + "github.com/omarluq/librecode/internal/tool" +) + +const ( + executeOuterCallID = "outer" + executeTestToolName = "echo" +) + +type executeTestTool struct{} + +func (executor *executeTestTool) Definition() tool.Definition { + return tool.Definition{ + Schema: mustToolSchema( + `{"type":"object","additionalProperties":false,"properties":` + + `{"text":{"type":"string"}},"required":["text"]}`, + ), + Name: executeTestToolName, Label: "Echo", Description: "Echo supplied text", + PromptSnippet: "Return supplied text", + PromptGuidelines: []string{"Pass text."}, ReadOnly: true, + } +} + +func (executor *executeTestTool) Execute(_ context.Context, input tool.Arguments) (tool.Result, error) { + var args struct { + Text string `json:"text"` + } + if err := input.Decode(&args); err != nil { + return tool.Result{}, oops.In("assistant").Code("execute_test_echo_input"). + Wrapf(err, "decode echo input") + } + + return tool.TextResult(args.Text, map[string]any{"length": len(args.Text)}), nil +} + +func TestExecuteToolUsesCompletedPromptRegistry(t *testing.T) { + t.Parallel() + + registry, err := tool.NewRegistryWithTools(t.TempDir(), nil) + require.NoError(t, err) + + echo := new(executeTestTool) + require.NoError(t, registry.Register(echo)) + execute := newExecuteTool(registry, nil) + require.NoError(t, registry.Register(execute)) + + tests := []struct { + check func(*testing.T, tool.Result) + name string + source string + }{ + { + name: "search", + source: `import "tools"; tools.Search("echo")`, + check: func(t *testing.T, result tool.Result) { + t.Helper() + assert.Contains(t, result.Text(), `"name":"echo"`) + assert.NotContains(t, result.Text(), `"name":"execute"`) + }, + }, + { + name: "describe includes schema", + source: `import "tools"; tools.Describe("echo")`, + check: func(t *testing.T, result tool.Result) { + t.Helper() + assert.Contains(t, result.Text(), `"required":["text"]`) + }, + }, + { + name: "multiple calls", + source: `import "tools" +first := tools.Call("echo", map[string]interface{}{"text": "one"}) +second := tools.Call("echo", map[string]interface{}{"text": "two"}) +[]interface{}{first, second}`, + check: func(t *testing.T, result tool.Result) { + t.Helper() + assert.Contains(t, result.Text(), `"text":"one"`) + assert.Contains(t, result.Text(), `"length":3`) + assert.NotNil(t, result.Details[executeResultValueKey]) + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + result, executeErr := execute.Execute(t.Context(), executeArguments(t, test.source)) + require.NoError(t, executeErr) + test.check(t, result) + }) + } +} + +func TestExecuteToolRejectsRecursionAndInvalidNestedInput(t *testing.T) { + t.Parallel() + + registry, err := tool.NewRegistryWithTools(t.TempDir(), nil) + require.NoError(t, err) + + execute := newExecuteTool(registry, nil) + require.NoError(t, registry.Register(execute)) + + tests := []struct { + name string + source string + want string + }{ + { + name: executeCallMethod, source: `import "tools"; tools.Call("execute", map[string]interface{}{})`, + want: "execute cannot call itself", + }, + {name: "describe", source: `import "tools"; tools.Describe("execute")`, want: "null"}, + { + name: "unknown", source: `import "tools"; tools.Call("missing", map[string]interface{}{})`, + want: "unknown tool", + }, + {name: "non object input", source: `import "tools"; tools.Call("missing", "bad")`, want: "cannot unmarshal"}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + result, executeErr := execute.Execute(t.Context(), executeArguments(t, test.source)) + require.NoError(t, executeErr) + assert.Contains(t, result.Text(), test.want) + }) + } +} + +func TestExecuteNestedCallsUseSharedInvocationBoundary(t *testing.T) { + t.Parallel() + + runtime := newToolExecutorTestRuntime(nil) + registry, err := tool.NewRegistryWithTools(t.TempDir(), nil) + require.NoError(t, err) + require.NoError(t, registry.Register(new(executeTestTool))) + + invoke := func( + ctx context.Context, + name string, + arguments tool.Arguments, + argumentsJSON string, + ) (tool.Result, ToolEvent) { + return runtime.invokeNestedTool(ctx, registry, name, arguments, argumentsJSON) + } + execute := newExecuteTool(registry, invoke) + require.NoError(t, registry.Register(execute)) + + events := []StreamEvent{} + outer := runtime.executeProviderToolCall( + t.Context(), registry, + &ToolCall{ + Metadata: nil, Arguments: executeArguments(t, `import "tools" + tools.Call("echo", map[string]interface{}{"text": "one"}) + tools.Call("echo", map[string]interface{}{"text": "two"})`), + ID: executeOuterCallID, Name: string(executeToolName), ArgumentsJSON: `{}`, + }, + func(event StreamEvent) { events = append(events, event) }, + ) + + require.False(t, outer.IsError) + require.Len(t, events, 6) + assert.Equal(t, StreamEventToolStart, events[0].Kind) + assert.Equal(t, executeOuterCallID, events[0].ToolCallEvent.ID) + assert.Equal(t, "echo", events[1].ToolCallEvent.Name) + assert.Equal(t, "outer/1", events[1].ToolCallEvent.ID) + assert.Equal(t, executeOuterCallID, events[1].ToolCallEvent.ParentCallID) + assert.Equal(t, 1, events[1].ToolCallEvent.Sequence) + assert.Equal(t, "one", events[2].ToolEvent.Result) + assert.Equal(t, events[1].ToolCallEvent.ID, events[2].ToolEvent.CallID) + assert.Equal(t, executeOuterCallID, events[2].ToolEvent.ParentCallID) + assert.Equal(t, 1, events[2].ToolEvent.Sequence) + assert.Equal(t, "echo", events[3].ToolCallEvent.Name) + assert.Equal(t, "outer/2", events[3].ToolCallEvent.ID) + assert.Equal(t, executeOuterCallID, events[3].ToolCallEvent.ParentCallID) + assert.Equal(t, 2, events[3].ToolCallEvent.Sequence) + assert.Equal(t, "two", events[4].ToolEvent.Result) + assert.Equal(t, events[3].ToolCallEvent.ID, events[4].ToolEvent.CallID) + assert.Equal(t, executeOuterCallID, events[4].ToolEvent.ParentCallID) + assert.Equal(t, 2, events[4].ToolEvent.Sequence) + assert.Equal(t, StreamEventToolResult, events[5].Kind) + assert.Equal(t, executeOuterCallID, events[5].ToolEvent.CallID) +} + +func TestExecuteToolHardCancelsHungProgram(t *testing.T) { + t.Parallel() + + registry, err := tool.NewRegistryWithTools(t.TempDir(), nil) + require.NoError(t, err) + + execute := newExecuteTool(registry, nil) + + ctx, cancel := context.WithTimeout(t.Context(), 300*time.Millisecond) + defer cancel() + + started := time.Now() + _, err = execute.Execute(ctx, executeArguments(t, `for {}; 1`)) + require.Error(t, err) + require.ErrorIs(t, err, context.DeadlineExceeded) + assert.Less(t, time.Since(started), 3*time.Second) +} + +func TestExecuteNestedCallsPreserveMixedResultsAndLifecycleOrder(t *testing.T) { + t.Parallel() + + recorder := &executeLifecycleRecorder{runtimeExtensions: nil, events: nil, mu: sync.Mutex{}} + runtime := newToolExecutorTestRuntime(recorder) + registry, err := tool.NewRegistryWithTools(t.TempDir(), nil) + require.NoError(t, err) + require.NoError(t, registry.Register(new(executeTestTool))) + execute := newExecuteTool(registry, func( + ctx context.Context, + name string, + arguments tool.Arguments, + argumentsJSON string, + ) (tool.Result, ToolEvent) { + return runtime.invokeNestedTool(ctx, registry, name, arguments, argumentsJSON) + }) + require.NoError(t, registry.Register(execute)) + + events := []StreamEvent{} + outer := runtime.executeProviderToolCall(t.Context(), registry, &ToolCall{ + Metadata: nil, + Arguments: executeArguments(t, `import "tools" +[]interface{}{ + tools.Call("echo", map[string]interface{}{"text": "ok"}), + tools.Call("missing", map[string]interface{}{}), +}`), + ID: executeOuterCallID, Name: string(executeToolName), ArgumentsJSON: `{}`, + }, func(event StreamEvent) { events = append(events, event) }) + + require.False(t, outer.IsError) + require.Len(t, events, 6) + assert.Equal(t, "ok", events[2].ToolEvent.Result) + assert.False(t, events[2].ToolEvent.IsError) + assert.Contains(t, events[4].ToolEvent.Error, "unknown tool") + assert.True(t, events[4].ToolEvent.IsError) + assert.Equal(t, []string{ + string(extension.LifecycleToolCall), + string(extension.LifecycleToolCall), + string(extension.LifecycleToolResult), + string(extension.LifecycleToolCall), + string(extension.LifecycleToolResult), + string(extension.LifecycleToolError), + string(extension.LifecycleToolResult), + }, recorder.eventNames()) +} + +func TestExecuteNestedCallCancellationEmitsCompletedBoundaries(t *testing.T) { + t.Parallel() + + runtime := newToolExecutorTestRuntime(nil) + registry, err := tool.NewRegistryWithTools(t.TempDir(), nil) + require.NoError(t, err) + + blocking := &executeBlockingTool{started: make(chan struct{})} + require.NoError(t, registry.Register(blocking)) + execute := newExecuteTool(registry, func( + ctx context.Context, + name string, + arguments tool.Arguments, + argumentsJSON string, + ) (tool.Result, ToolEvent) { + return runtime.invokeNestedTool(ctx, registry, name, arguments, argumentsJSON) + }) + require.NoError(t, registry.Register(execute)) + + ctx, cancel := context.WithCancel(t.Context()) + events := []StreamEvent{} + + done := make(chan ToolEvent, 1) + arguments := executeArguments(t, `import "tools"; tools.Call("block", map[string]interface{}{})`) + + go func() { + done <- runtime.executeProviderToolCall(ctx, registry, &ToolCall{ + Metadata: nil, + Arguments: arguments, + ID: executeOuterCallID, Name: string(executeToolName), ArgumentsJSON: `{}`, + }, func(event StreamEvent) { events = append(events, event) }) + }() + + <-blocking.started + cancel() + + outer := <-done + + assert.True(t, outer.IsError) + require.Len(t, events, 4) + assert.Equal(t, []StreamEventKind{ + StreamEventToolStart, StreamEventToolStart, StreamEventToolResult, StreamEventToolResult, + }, []StreamEventKind{events[0].Kind, events[1].Kind, events[2].Kind, events[3].Kind}) + assert.Equal(t, executeOuterCallID, events[2].ToolEvent.ParentCallID) + assert.True(t, events[2].ToolEvent.IsError) +} + +func TestExecuteResultTextRejectsOversizedResult(t *testing.T) { + t.Parallel() + + _, err := executeResultText(mvmhost.Result{ + Value: strings.Repeat("x", defaultExecuteResultLimit), ValueKind: "", Stdout: "", Stderr: "", + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "execute result") + assert.Contains(t, err.Error(), "limit") +} + +func TestExecuteToolValidationAndWorkerErrors(t *testing.T) { + t.Parallel() + + missingRegistry := newExecuteTool(nil, nil) + _, err := missingRegistry.Execute(t.Context(), tool.EmptyArguments()) + require.ErrorContains(t, err, "registry is not configured") + + registry, registryErr := tool.NewRegistryWithTools(t.TempDir(), nil) + require.NoError(t, registryErr) + + execute := newExecuteTool(registry, nil) + + invalidInput, inputErr := tool.ArgumentsFromRaw([]byte(`{"source":42}`)) + require.NoError(t, inputErr) + _, err = execute.Execute(t.Context(), invalidInput) + require.ErrorContains(t, err, "decode execute input") + + _, err = execute.handleWorkerMessage(t.Context(), executeTestWorkerMessage("unknown", "", nil)) + require.ErrorContains(t, err, "unknown execute worker RPC method") + + _, err = execute.handleWorkerMessage( + t.Context(), + executeTestWorkerMessage(executeCallMethod, executeTestToolName, json.RawMessage(`[]`)), + ) + require.ErrorContains(t, err, "decode nested tool input") +} + +func executeTestWorkerMessage(method, name string, input json.RawMessage) *executeworker.Message { + return &executeworker.Message{ + Stderr: "", Source: "", Method: method, Mode: "", Name: name, Query: "", Stdout: "", Type: "", + Error: "", ErrorKind: "", ValueKind: "", Input: input, Value: nil, Arguments: nil, ID: 0, + ExitCode: 0, + } +} + +func TestExecuteCallHandlesProtectedToolsAndInvocationError(t *testing.T) { + t.Parallel() + + registry, err := tool.NewRegistryWithTools(t.TempDir(), nil) + require.NoError(t, err) + + execute := newExecuteTool(registry, func( + context.Context, + string, + tool.Arguments, + string, + ) (tool.Result, ToolEvent) { + return tool.Result{Details: nil, Content: nil}, ToolEvent{ + CallID: "", ParentCallID: "", Name: executeTestToolName, ArgumentsJSON: "", DetailsJSON: "", + Result: "", Error: "nested failure", Sequence: 0, IsError: true, + } + }) + + workflowResult := execute.call(t.Context(), string(workflowToolName), json.RawMessage(`{}`)) + assert.True(t, workflowResult.IsError) + assert.Equal(t, "execute cannot call workflow", workflowResult.Error) + + invalidResult := execute.call(t.Context(), executeTestToolName, json.RawMessage(`[]`)) + assert.True(t, invalidResult.IsError) + assert.Contains(t, invalidResult.Error, "decode tool arguments") + + invocationResult := execute.call(t.Context(), executeTestToolName, json.RawMessage(`{}`)) + assert.True(t, invocationResult.IsError) + assert.Equal(t, "nested failure", invocationResult.Error) +} + +func TestExecuteResultTextVariants(t *testing.T) { + t.Parallel() + + text, err := executeResultText(mvmhost.Result{ + Value: nil, ValueKind: "", Stdout: "printed", Stderr: "", + }) + require.NoError(t, err) + assert.Equal(t, "printed", text) + + text, err = executeResultText(mvmhost.Result{Value: nil, ValueKind: "", Stdout: "", Stderr: ""}) + require.NoError(t, err) + assert.Equal(t, "null", text) + + text, err = executeResultText(mvmhost.Result{ + Value: map[string]any{"ok": true}, ValueKind: "", Stdout: "", Stderr: "", + }) + require.NoError(t, err) + assert.JSONEq(t, `{"ok":true}`, text) + + _, err = executeResultText(mvmhost.Result{ + Value: make(chan struct{}), ValueKind: "", Stdout: "", Stderr: "", + }) + require.ErrorContains(t, err, "encode execute result") + + definition := tool.Definition{ + Schema: tool.EmptySchema(), Name: "empty", Label: "", Description: "", PromptSnippet: "", + PromptGuidelines: nil, ReadOnly: false, + } + assert.Nil(t, executeDefinitionMap(&definition)["schema"]) +} + +type executeBlockingTool struct { + started chan struct{} +} + +func (executor *executeBlockingTool) Definition() tool.Definition { + return tool.Definition{ + Schema: mustToolSchema(`{"type":"object","additionalProperties":false}`), + Name: "block", Label: "Block", Description: "Block until canceled", PromptSnippet: "Block", + PromptGuidelines: nil, ReadOnly: true, + } +} + +func (executor *executeBlockingTool) Execute(ctx context.Context, _ tool.Arguments) (tool.Result, error) { + close(executor.started) + <-ctx.Done() + + return tool.Result{}, oops.In("assistant").Code("blocking_tool_canceled").Wrapf(ctx.Err(), "wait for cancellation") +} + +type executeLifecycleRecorder struct { + runtimeExtensions + events []string + mu sync.Mutex +} + +func (recorder *executeLifecycleRecorder) DispatchLifecycle( + _ context.Context, + event extension.LifecycleEvent, +) (extension.LifecycleDispatchResult, error) { + recorder.mu.Lock() + defer recorder.mu.Unlock() + + recorder.events = append(recorder.events, string(event.Name)) + + return emptyTestLifecycleDispatchResult(event), nil +} + +func (recorder *executeLifecycleRecorder) eventNames() []string { + recorder.mu.Lock() + defer recorder.mu.Unlock() + + return append([]string(nil), recorder.events...) +} + +func TestPromptRegistryRegistersExecuteAfterPromptLocalTools(t *testing.T) { + t.Parallel() + + runtime := new(Runtime) + runtime.profile = topLevelExecutionProfile() + registry, err := runtime.promptToolRegistry(t.Context(), t.TempDir(), "owner") + require.NoError(t, err) + + definitions := registry.Definitions() + found := false + + for _, definition := range definitions { + if definition.Name == executeToolName { + found = true + + assert.False(t, definition.ReadOnly) + } + } + + assert.True(t, found) +} + +func executeArguments(t *testing.T, source string) tool.Arguments { + t.Helper() + + encoded, err := json.Marshal(map[string]string{"source": source}) + require.NoError(t, err) + arguments, err := tool.ArgumentsFromRaw(encoded) + require.NoError(t, err) + + return arguments +} + +var _ tool.Executor = (*executeTestTool)(nil) diff --git a/internal/assistant/lifecycle.go b/internal/assistant/lifecycle.go index 8590061b..5e5deeaf 100644 --- a/internal/assistant/lifecycle.go +++ b/internal/assistant/lifecycle.go @@ -201,7 +201,7 @@ func (runtime *Runtime) dispatchToolCallLifecycle(ctx context.Context, call *Too return nil } - payload := lifecyclepayload.ToolCallPayload(lifecycleToolCall(*call)) + payload := lifecyclepayload.ToolCallPayload(lifecycleToolCall(call)) result, err := runtime.dispatchLifecycle(ctx, extension.LifecycleToolCall, payload) if err != nil { @@ -291,22 +291,27 @@ func lifecyclePromptRequest(request *PromptRequest) *lifecyclepayload.PromptRequ } } -func lifecycleToolCall(call ToolCallEvent) lifecyclepayload.ToolCall { - return lifecyclepayload.ToolCall{ +func lifecycleToolCall(call *ToolCallEvent) *lifecyclepayload.ToolCall { + return &lifecyclepayload.ToolCall{ Arguments: call.Arguments, ID: call.ID, + ParentCallID: call.ParentCallID, Name: call.Name, ArgumentsJSON: call.ArgumentsJSON, + Sequence: call.Sequence, } } func lifecycleToolResult(event *ToolEvent) *lifecyclepayload.ToolResult { return &lifecyclepayload.ToolResult{ + CallID: event.CallID, + ParentCallID: event.ParentCallID, Name: event.Name, ArgumentsJSON: event.ArgumentsJSON, DetailsJSON: event.DetailsJSON, Result: event.Result, Error: event.Error, + Sequence: event.Sequence, IsError: event.IsError, } } diff --git a/internal/assistant/lifecyclepayload/lifecyclepayload.go b/internal/assistant/lifecyclepayload/lifecyclepayload.go index f9e1bfb9..ccfca71c 100644 --- a/internal/assistant/lifecyclepayload/lifecyclepayload.go +++ b/internal/assistant/lifecyclepayload/lifecyclepayload.go @@ -110,17 +110,22 @@ type ProviderError struct { type ToolCall struct { ArgumentsJSON string ID string + ParentCallID string Name string Arguments tool.Arguments + Sequence int } // ToolResult describes a tool-result lifecycle payload. type ToolResult struct { + CallID string + ParentCallID string Name string ArgumentsJSON string DetailsJSON string Result string Error string + Sequence int IsError bool } @@ -339,9 +344,15 @@ func ProviderErrorPayload(providerErr *ProviderError) map[string]any { } // ToolCallPayload builds tool-call lifecycle payloads. -func ToolCallPayload(call ToolCall) map[string]any { +func ToolCallPayload(call *ToolCall) map[string]any { + if call == nil { + return map[string]any{} + } + return map[string]any{ "call_id": call.ID, + "parent_call_id": call.ParentCallID, + "sequence": call.Sequence, ToolNameKey: call.Name, "arguments_json": call.ArgumentsJSON, "arguments": call.Arguments, @@ -355,6 +366,9 @@ func ToolResultPayload(event *ToolResult) map[string]any { } return map[string]any{ + "call_id": event.CallID, + "parent_call_id": event.ParentCallID, + "sequence": event.Sequence, ToolNameKey: event.Name, "arguments_json": event.ArgumentsJSON, "details_json": event.DetailsJSON, diff --git a/internal/assistant/lifecyclepayload/lifecyclepayload_test.go b/internal/assistant/lifecyclepayload/lifecyclepayload_test.go index 11bb4a61..4391cd6d 100644 --- a/internal/assistant/lifecyclepayload/lifecyclepayload_test.go +++ b/internal/assistant/lifecyclepayload/lifecyclepayload_test.go @@ -18,6 +18,7 @@ import ( const ( lifecycleTestSessionID = "session-1" + lifecycleTestParentID = "parent-1" lifecycleTestAPI = "openai-responses" lifecycleTestModel = "model-1" lifecycleTestProvider = "openai" @@ -26,7 +27,7 @@ const ( func TestPromptAndTurnPayloads(t *testing.T) { t.Parallel() - parentID := "parent-1" + parentID := lifecycleTestParentID prompt := lifecyclepayload.Prompt(&lifecyclepayload.PromptRequest{ ParentEntryID: &parentID, CWD: "/work", @@ -155,23 +156,33 @@ func TestProviderAndToolPayloads(t *testing.T) { assert.Equal(t, true, providerErr[lifecyclepayload.ProviderBodyTruncatedKey]) assert.Equal(t, map[string]any{"input_count": 2}, providerErr[lifecyclepayload.ProviderRequestShapeKey]) - toolCall := lifecyclepayload.ToolCallPayload(lifecyclepayload.ToolCall{ + toolCall := lifecyclepayload.ToolCallPayload(&lifecyclepayload.ToolCall{ Arguments: testutil.ToolArguments(map[string]any{"path": "README.md"}), ID: "call-1", + ParentCallID: lifecycleTestParentID, Name: "read", ArgumentsJSON: `{"path":"README.md"}`, + Sequence: 2, }) assert.Equal(t, "call-1", toolCall["call_id"]) + assert.Equal(t, lifecycleTestParentID, toolCall["parent_call_id"]) + assert.Equal(t, 2, toolCall["sequence"]) assert.Equal(t, "read", toolCall[lifecyclepayload.ToolNameKey]) toolResult := lifecyclepayload.ToolResultPayload(&lifecyclepayload.ToolResult{ + CallID: "call-1", + ParentCallID: lifecycleTestParentID, Name: "read", ArgumentsJSON: "{}", DetailsJSON: "{}", Result: "ok", Error: "failed", + Sequence: 2, IsError: true, }) + assert.Equal(t, "call-1", toolResult["call_id"]) + assert.Equal(t, lifecycleTestParentID, toolResult["parent_call_id"]) + assert.Equal(t, 2, toolResult["sequence"]) assert.Equal(t, true, toolResult["is_error"]) assert.Equal(t, "failed", toolResult[lifecyclepayload.ToolErrorKey]) } diff --git a/internal/assistant/llm_conversion.go b/internal/assistant/llm_conversion.go index a9303958..acbc2970 100644 --- a/internal/assistant/llm_conversion.go +++ b/internal/assistant/llm_conversion.go @@ -179,8 +179,8 @@ func llmToolResultFromEvent(event *ToolEvent) *llm.ToolResult { } return &llm.ToolResult{ - Metadata: nil, - ToolCallID: "", + Metadata: toolIdentityMetadata(event.ParentCallID, event.Sequence), + ToolCallID: event.CallID, ArgumentsJSON: event.ArgumentsJSON, Name: event.Name, Error: event.Error, diff --git a/internal/assistant/llm_conversion_internal_test.go b/internal/assistant/llm_conversion_internal_test.go index 7f165d48..a13e4a2c 100644 --- a/internal/assistant/llm_conversion_internal_test.go +++ b/internal/assistant/llm_conversion_internal_test.go @@ -119,21 +119,28 @@ func TestLLMResponseFromCompletionResultConvertsContentAndUsage(t *testing.T) { FinishReason: llm.FinishReasonLength, Text: "final answer", Thinking: []string{" thought ", " "}, - ToolEvents: []ToolEvent{{ - Name: jsonReadToolName, - ArgumentsJSON: `{"path":"README.md"}`, - DetailsJSON: "", - Result: "read output", - Error: "", - IsError: false, - }, { - Name: "bash", - ArgumentsJSON: `{"command":"false"}`, - DetailsJSON: "", - Result: "exit status 1", - Error: "exit status 1", - IsError: true, - }}, + ToolEvents: []ToolEvent{ + { + CallID: "conversion-call-1", + ParentCallID: "conversion-parent", + Sequence: 1, + Name: jsonReadToolName, + ArgumentsJSON: `{"path":"README.md"}`, + DetailsJSON: "", + Result: "read output", + Error: "", + IsError: false, + }, { + CallID: "conversion-call-2", + ParentCallID: "conversion-parent", + Sequence: 2, + Name: "bash", + ArgumentsJSON: `{"command":"false"}`, + DetailsJSON: "", + Result: "exit status 1", + Error: "exit status 1", + IsError: true, + }}, Usage: model.TokenUsage{ Breakdown: map[string]int{contextwindow.BreakdownHistory: 10}, TopContributors: nil, @@ -155,8 +162,14 @@ func TestLLMResponseFromCompletionResultConvertsContentAndUsage(t *testing.T) { assert.Equal(t, "final answer", converted.Content[1].Text) assert.Equal(t, llm.PartToolResult, converted.Content[2].Type) assert.False(t, converted.Content[2].ToolResult.IsError) + assert.Equal(t, "conversion-call-1", converted.Content[2].ToolResult.ToolCallID) + assert.Equal(t, "conversion-parent", converted.Content[2].ToolResult.Metadata[toolParentCallIDMetadataKey]) + assert.Equal(t, 1, converted.Content[2].ToolResult.Metadata[toolSequenceMetadataKey]) assert.Equal(t, llm.PartToolResult, converted.Content[3].Type) assert.True(t, converted.Content[3].ToolResult.IsError) + assert.Equal(t, "conversion-call-2", converted.Content[3].ToolResult.ToolCallID) + assert.Equal(t, "conversion-parent", converted.Content[3].ToolResult.Metadata[toolParentCallIDMetadataKey]) + assert.Equal(t, 2, converted.Content[3].ToolResult.Metadata[toolSequenceMetadataKey]) assert.Equal(t, "exit status 1", converted.Content[3].ToolResult.Error) } diff --git a/internal/assistant/provider_hook_test_helpers_internal_test.go b/internal/assistant/provider_hook_test_helpers_internal_test.go index c54fd73f..5aa04726 100644 --- a/internal/assistant/provider_hook_test_helpers_internal_test.go +++ b/internal/assistant/provider_hook_test_helpers_internal_test.go @@ -42,18 +42,19 @@ func newRuntimeFromDeps(setup func(*runtimeDeps)) *Runtime { } return &Runtime{ - cfg: deps.Config, - sessions: deps.Sessions, - extensions: deps.Extensions, - cache: deps.Cache, - models: deps.Models, - client: client, - logger: deps.Logger, - skillsCache: nil, - toolSchemaCache: newToolSchemaCache(), - agents: nil, - agentTasks: nil, - profile: topLevelExecutionProfile(), + cfg: deps.Config, + sessions: deps.Sessions, + extensions: deps.Extensions, + cache: deps.Cache, + models: deps.Models, + client: client, + logger: deps.Logger, + skillsCache: nil, + toolSchemaCache: newToolSchemaCache(), + agents: nil, + agentTasks: nil, + workflowSubmitter: nil, + profile: topLevelExecutionProfile(), } } diff --git a/internal/assistant/run_metrics.go b/internal/assistant/run_metrics.go new file mode 100644 index 00000000..edf669e0 --- /dev/null +++ b/internal/assistant/run_metrics.go @@ -0,0 +1,218 @@ +package assistant + +import ( + "context" + "sync" + + "github.com/omarluq/librecode/internal/model" +) + +// ToolStrategy controls which provider-facing tool surfaces are available. +type ToolStrategy string + +const ( + // ToolStrategyHybrid exposes direct tools and the execute code-mode tool. + ToolStrategyHybrid ToolStrategy = "hybrid" + // ToolStrategyDirect exposes direct tools without execute. + ToolStrategyDirect ToolStrategy = "direct" +) + +type toolStrategyContextKey struct{} +type runMetricsContextKey struct{} + +// WithToolStrategy returns a context that selects a provider-facing tool strategy. +func WithToolStrategy(ctx context.Context, strategy ToolStrategy) context.Context { + return context.WithValue(ctx, toolStrategyContextKey{}, strategy) +} + +func toolStrategyFromContext(ctx context.Context) ToolStrategy { + strategy, ok := ctx.Value(toolStrategyContextKey{}).(ToolStrategy) + if !ok || strategy == "" { + return ToolStrategyHybrid + } + + return strategy +} + +// RunMetricsSnapshot is a stable point-in-time view of prompt execution metrics. +type RunMetricsSnapshot struct { + ProviderRoundTrips int + InputTokens int + OutputTokens int + ToolCalls int + NestedToolCalls int + TraceComplete bool +} + +// RunMetrics collects request-local provider, usage, and tool-trace observations. +type RunMetrics struct { + starts map[string]ToolCallEvent + results map[string]struct{} + providerRoundTrips int + inputTokens int + outputTokens int + toolCalls int + nestedToolCalls int + traceValid bool + initialized bool + mu sync.Mutex +} + +// WithRunMetrics returns a context that records prompt execution metrics. +func WithRunMetrics(ctx context.Context, metrics *RunMetrics) context.Context { + if metrics == nil { + return ctx + } + + return context.WithValue(ctx, runMetricsContextKey{}, metrics) +} + +func runMetricsFromContext(ctx context.Context) *RunMetrics { + metrics, ok := ctx.Value(runMetricsContextKey{}).(*RunMetrics) + if !ok { + return nil + } + + return metrics +} + +func observeProviderRoundTrip(ctx context.Context) { + metrics := runMetricsFromContext(ctx) + if metrics == nil { + return + } + + metrics.mu.Lock() + metrics.initialize() + metrics.providerRoundTrips++ + metrics.mu.Unlock() +} + +func observeProviderUsage(ctx context.Context, usage model.TokenUsage) { + metrics := runMetricsFromContext(ctx) + if metrics == nil { + return + } + + metrics.mu.Lock() + metrics.initialize() + metrics.inputTokens = usage.InputTokens + metrics.outputTokens = usage.OutputTokens + metrics.mu.Unlock() +} + +// ObserveStreamEvent records tool starts and results for trace accounting. +func (metrics *RunMetrics) ObserveStreamEvent(event StreamEvent) { + if metrics == nil { + return + } + + metrics.mu.Lock() + defer metrics.mu.Unlock() + + metrics.initialize() + + switch event.Kind { + case StreamEventToolStart: + metrics.observeToolStart(event.ToolCallEvent) + case StreamEventToolResult: + metrics.observeToolResult(event.ToolEvent) + case StreamEventUsage, StreamEventUsageSnapshot: + metrics.observeUsage(event.Usage) + case StreamEventTextDelta, + StreamEventThinkingDelta, + StreamEventSkillLoaded, + StreamEventContextCompaction, + StreamEventContextCompactionStart, + StreamEventContextCompactionDone, + StreamEventContextCompactionError, + StreamEventUnknown: + } +} + +func (metrics *RunMetrics) observeUsage(usage *model.TokenUsage) { + if usage == nil { + return + } + + metrics.inputTokens = usage.InputTokens + metrics.outputTokens = usage.OutputTokens +} + +func (metrics *RunMetrics) observeToolStart(event *ToolCallEvent) { + if event == nil || event.ID == "" { + metrics.traceValid = false + + return + } + + if _, exists := metrics.starts[event.ID]; exists { + metrics.traceValid = false + + return + } + + metrics.starts[event.ID] = *event +} + +func (metrics *RunMetrics) observeToolResult(event *ToolEvent) { + metrics.toolCalls++ + if event == nil || event.CallID == "" { + metrics.traceValid = false + + return + } + + if event.ParentCallID != "" { + metrics.nestedToolCalls++ + } + + start, exists := metrics.starts[event.CallID] + + _, duplicate := metrics.results[event.CallID] + if !exists || duplicate || start.ParentCallID != event.ParentCallID || start.Sequence != event.Sequence { + metrics.traceValid = false + } + + metrics.results[event.CallID] = struct{}{} +} + +// Snapshot returns a concurrency-safe copy of the collected metrics. +func (metrics *RunMetrics) Snapshot() RunMetricsSnapshot { + if metrics == nil { + return RunMetricsSnapshot{ + ProviderRoundTrips: 0, InputTokens: 0, OutputTokens: 0, + ToolCalls: 0, NestedToolCalls: 0, TraceComplete: true, + } + } + + metrics.mu.Lock() + defer metrics.mu.Unlock() + + metrics.initialize() + + return RunMetricsSnapshot{ + ProviderRoundTrips: metrics.providerRoundTrips, + InputTokens: metrics.inputTokens, + OutputTokens: metrics.outputTokens, + ToolCalls: metrics.toolCalls, + NestedToolCalls: metrics.nestedToolCalls, + TraceComplete: metrics.traceValid && len(metrics.starts) == len(metrics.results), + } +} + +// ProviderRoundTrips returns the observed provider request count. +func (metrics *RunMetrics) ProviderRoundTrips() int { + return metrics.Snapshot().ProviderRoundTrips +} + +func (metrics *RunMetrics) initialize() { + if metrics.initialized { + return + } + + metrics.starts = make(map[string]ToolCallEvent) + metrics.results = make(map[string]struct{}) + metrics.traceValid = true + metrics.initialized = true +} diff --git a/internal/assistant/run_metrics_internal_test.go b/internal/assistant/run_metrics_internal_test.go new file mode 100644 index 00000000..3f002a55 --- /dev/null +++ b/internal/assistant/run_metrics_internal_test.go @@ -0,0 +1,186 @@ +package assistant + +import ( + "sync" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/omarluq/librecode/internal/model" + "github.com/omarluq/librecode/internal/tool" +) + +const metricsMalformedCall = "metrics-malformed-call" + +func TestRunMetricsCollectsUsageAndNestedTrace(t *testing.T) { + t.Parallel() + + const ( + outerCallID = "metrics-outer" + nestedCallID = "metrics-outer/1" + ) + + metrics := new(RunMetrics) + ctx := WithRunMetrics(t.Context(), metrics) + observeProviderRoundTrip(ctx) + observeProviderRoundTrip(ctx) + observeProviderUsage(ctx, metricsTokenUsage(13, 5)) + + metrics.ObserveStreamEvent(metricsStart(outerCallID, "", 0)) + metrics.ObserveStreamEvent(metricsStart(nestedCallID, outerCallID, 1)) + metrics.ObserveStreamEvent(metricsResult(nestedCallID, outerCallID, 1)) + metrics.ObserveStreamEvent(metricsResult(outerCallID, "", 0)) + + assert.Equal(t, RunMetricsSnapshot{ + ProviderRoundTrips: 2, InputTokens: 13, OutputTokens: 5, + ToolCalls: 2, NestedToolCalls: 1, TraceComplete: true, + }, metrics.Snapshot()) +} + +func TestRunMetricsDetectsIncompleteTrace(t *testing.T) { + t.Parallel() + + metrics := new(RunMetrics) + metrics.ObserveStreamEvent(metricsStart("metrics-incomplete", "", 0)) + + assert.False(t, metrics.Snapshot().TraceComplete) +} + +func TestRunMetricsContextDefaultsAndNilSafety(t *testing.T) { + t.Parallel() + + ctx := t.Context() + assert.Equal(t, ToolStrategyHybrid, toolStrategyFromContext(ctx)) + assert.Equal(t, ToolStrategyHybrid, toolStrategyFromContext(WithToolStrategy(ctx, ""))) + assert.Equal(t, ToolStrategyDirect, toolStrategyFromContext(WithToolStrategy(ctx, ToolStrategyDirect))) + assert.Same(t, ctx, WithRunMetrics(ctx, nil)) + assert.Nil(t, runMetricsFromContext(ctx)) + + observeProviderRoundTrip(ctx) + observeProviderUsage(ctx, metricsTokenUsage(9, 4)) + + var nilMetrics *RunMetrics + + nilMetrics.ObserveStreamEvent(metricsEvent(StreamEventToolStart, nil, nil, nil)) + assert.Equal(t, RunMetricsSnapshot{ + ProviderRoundTrips: 0, InputTokens: 0, OutputTokens: 0, + ToolCalls: 0, NestedToolCalls: 0, TraceComplete: true, + }, nilMetrics.Snapshot()) +} + +func TestRunMetricsStreamUsageAndIgnoredEvents(t *testing.T) { + t.Parallel() + + metrics := new(RunMetrics) + metrics.ObserveStreamEvent(metricsEvent(StreamEventUsage, nil, nil, nil)) + + usage := metricsTokenUsage(21, 8) + metrics.ObserveStreamEvent(metricsEvent(StreamEventUsageSnapshot, nil, nil, &usage)) + + for _, kind := range []StreamEventKind{ + StreamEventTextDelta, StreamEventThinkingDelta, StreamEventSkillLoaded, + StreamEventContextCompaction, StreamEventContextCompactionStart, + StreamEventContextCompactionDone, StreamEventContextCompactionError, StreamEventUnknown, + } { + metrics.ObserveStreamEvent(metricsEvent(kind, nil, nil, nil)) + } + + snapshot := metrics.Snapshot() + assert.Equal(t, 21, snapshot.InputTokens) + assert.Equal(t, 8, snapshot.OutputTokens) +} + +func TestRunMetricsRejectsMalformedToolTraces(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + events []StreamEvent + }{ + {name: "nil start", events: []StreamEvent{metricsEvent(StreamEventToolStart, nil, nil, nil)}}, + {name: "empty start ID", events: []StreamEvent{metricsStart("", "", 0)}}, + {name: "duplicate start", events: []StreamEvent{ + metricsStart(metricsMalformedCall, "", 0), metricsStart(metricsMalformedCall, "", 0), + }}, + {name: "nil result", events: []StreamEvent{metricsEvent(StreamEventToolResult, nil, nil, nil)}}, + {name: "empty result ID", events: []StreamEvent{metricsResult("", "", 0)}}, + {name: "result without start", events: []StreamEvent{metricsResult(metricsMalformedCall, "", 0)}}, + {name: "duplicate result", events: []StreamEvent{ + metricsStart(metricsMalformedCall, "", 0), + metricsResult(metricsMalformedCall, "", 0), + metricsResult(metricsMalformedCall, "", 0), + }}, + {name: "identity mismatch", events: []StreamEvent{ + metricsStart(metricsMalformedCall, "parent", 1), + metricsResult(metricsMalformedCall, "other", 2), + }}, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + metrics := new(RunMetrics) + for _, event := range testCase.events { + metrics.ObserveStreamEvent(event) + } + + assert.False(t, metrics.Snapshot().TraceComplete) + }) + } +} + +func metricsTokenUsage(input, output int) model.TokenUsage { + return model.TokenUsage{ + Breakdown: nil, TopContributors: nil, ContextWindow: 0, ContextTokens: 0, + InputTokens: input, OutputTokens: output, + } +} + +func metricsEvent( + kind StreamEventKind, + call *ToolCallEvent, + result *ToolEvent, + usage *model.TokenUsage, +) StreamEvent { + return StreamEvent{ToolCallEvent: call, ToolEvent: result, Usage: usage, Kind: kind, Text: ""} +} + +func metricsStart(id, parent string, sequence int) StreamEvent { + return metricsEvent(StreamEventToolStart, &ToolCallEvent{ + ArgumentsJSON: "", ID: id, ParentCallID: parent, Name: "", + Arguments: tool.EmptyArguments(), Sequence: sequence, + }, nil, nil) +} + +func metricsResult(id, parent string, sequence int) StreamEvent { + return metricsEvent(StreamEventToolResult, nil, &ToolEvent{ + CallID: id, ParentCallID: parent, Name: "", ArgumentsJSON: "", DetailsJSON: "", + Result: "", Error: "", Sequence: sequence, IsError: false, + }, nil) +} + +func TestRunMetricsConcurrentProviderObservations(t *testing.T) { + t.Parallel() + + const observations = 32 + + metrics := new(RunMetrics) + ctx := WithRunMetrics(t.Context(), metrics) + + var wait sync.WaitGroup + + wait.Add(observations) + + for range observations { + go func() { + defer wait.Done() + + observeProviderRoundTrip(ctx) + }() + } + + wait.Wait() + + assert.Equal(t, observations, metrics.ProviderRoundTrips()) +} diff --git a/internal/assistant/runtime.go b/internal/assistant/runtime.go index 1fb7ca02..8ca95601 100644 --- a/internal/assistant/runtime.go +++ b/internal/assistant/runtime.go @@ -27,18 +27,19 @@ type runtimeExtensions interface { // Runtime coordinates prompt handling and durable sessions. type Runtime struct { - cfg *config.Config - sessions *database.SessionRepository - extensions runtimeExtensions - cache *ResponseCache - models *model.Registry - client Completer - logger *slog.Logger - skillsCache *core.SkillsCache - toolSchemaCache *toolSchemaCache - agents *agent.Catalog - agentTasks AgentTaskController - profile ExecutionProfile + cfg *config.Config + sessions *database.SessionRepository + extensions runtimeExtensions + cache *ResponseCache + models *model.Registry + client Completer + logger *slog.Logger + skillsCache *core.SkillsCache + toolSchemaCache *toolSchemaCache + agents *agent.Catalog + agentTasks AgentTaskController + workflowSubmitter workflowSubmitter + profile ExecutionProfile } // PromptRequest contains one user prompt invocation. @@ -149,18 +150,19 @@ func NewRuntime(options *RuntimeOptions) *Runtime { } return &Runtime{ - cfg: options.Config, - sessions: options.Sessions, - extensions: options.Extensions, - cache: options.Cache, - models: options.Models, - client: client, - logger: options.Logger, - skillsCache: options.SkillsCache, - toolSchemaCache: newToolSchemaCache(), - agents: options.Agents, - agentTasks: nil, - profile: topLevelExecutionProfile(), + cfg: options.Config, + sessions: options.Sessions, + extensions: options.Extensions, + cache: options.Cache, + models: options.Models, + client: client, + logger: options.Logger, + skillsCache: options.SkillsCache, + toolSchemaCache: newToolSchemaCache(), + agents: options.Agents, + agentTasks: nil, + workflowSubmitter: nil, + profile: topLevelExecutionProfile(), } } @@ -170,6 +172,12 @@ func (runtime *Runtime) SetAgentTaskController(controller AgentTaskController) { runtime.agentTasks = controller } +// SetWorkflowSubmitter enables model-authored durable workflows. +// It must be called during application startup, before Prompt is used. +func (runtime *Runtime) SetWorkflowSubmitter(submitter workflowSubmitter) { + runtime.workflowSubmitter = submitter +} + // Prompt appends a user prompt and an assistant response to the selected session. func (runtime *Runtime) Prompt(ctx context.Context, request *PromptRequest) (response *PromptResponse, err error) { if request == nil { @@ -414,7 +422,7 @@ func (runtime *Runtime) WithExecutionProfile(profile *ExecutionProfile) *Runtime cfg: runtime.cfg, sessions: runtime.sessions, extensions: runtime.extensions, cache: runtime.cache, models: runtime.models, client: runtime.client, logger: runtime.logger, skillsCache: runtime.skillsCache, toolSchemaCache: runtime.toolSchemaCache, agents: runtime.agents, - agentTasks: runtime.agentTasks, profile: clonedProfile, + agentTasks: runtime.agentTasks, workflowSubmitter: runtime.workflowSubmitter, profile: clonedProfile, } } diff --git a/internal/assistant/runtime_events.go b/internal/assistant/runtime_events.go index c1ac7f14..fb0431ca 100644 --- a/internal/assistant/runtime_events.go +++ b/internal/assistant/runtime_events.go @@ -12,6 +12,8 @@ func (runtime *Runtime) emitProviderRequest(ctx context.Context, request *Comple return } + observeProviderRoundTrip(ctx) + payload := new(lifecyclepayload.ProviderRequest) payload.Headers = redactedHeaders(request.Auth.Headers) payload.API = request.Model.API @@ -44,6 +46,8 @@ func (runtime *Runtime) emitProviderResponse( return } + observeProviderUsage(ctx, result.Usage) + payload := lifecyclepayload.ProviderResponsePayload(&lifecyclepayload.ProviderResponse{ FinishReason: result.FinishReason, Usage: result.Usage, diff --git a/internal/assistant/runtime_lifecycle_test.go b/internal/assistant/runtime_lifecycle_test.go index b6885232..d4caea88 100644 --- a/internal/assistant/runtime_lifecycle_test.go +++ b/internal/assistant/runtime_lifecycle_test.go @@ -175,6 +175,9 @@ func TestRuntime_PromptEmitsSideEffectMessageAppendEvents(t *testing.T) { Thinking: []string{"reasoning"}, ToolEvents: []assistant.ToolEvent{ { + CallID: "", + ParentCallID: "", + Sequence: 0, Name: testToolName, ArgumentsJSON: testToolArgsJSON, DetailsJSON: "", diff --git a/internal/assistant/runtime_persist.go b/internal/assistant/runtime_persist.go index 1c901496..2eb4acb6 100644 --- a/internal/assistant/runtime_persist.go +++ b/internal/assistant/runtime_persist.go @@ -27,8 +27,11 @@ type partialPromptBlock struct { } type pendingToolCall struct { + CallID string + ParentCallID string Name string ArgumentsJSON string + Sequence int } type partialPromptMessage struct { @@ -37,10 +40,11 @@ type partialPromptMessage struct { } type partialPromptProgress struct { - forward func(StreamEvent) - blocks []partialPromptBlock - fallbackBlocks []partialPromptBlock - pendingTools []pendingToolCall + forward func(StreamEvent) + blocks []partialPromptBlock + fallbackBlocks []partialPromptBlock + pendingTools []pendingToolCall + completedNestedTools []ToolEvent } func (runtime *Runtime) appendAssistantSideEffects( @@ -72,11 +76,12 @@ func (runtime *Runtime) appendAssistantSideEffects( parentID = &entry.ID } - for _, event := range bundle.ToolEvents { + for index := range bundle.ToolEvents { + event := &bundle.ToolEvents[index] message := database.MessageEntity{ Timestamp: time.Now().UTC(), Role: database.RoleToolResult, - Content: formatToolEvent(&event), + Content: formatToolEvent(event), Provider: runtime.cfg.Assistant.Provider, Model: runtime.cfg.Assistant.Model, } @@ -122,15 +127,138 @@ func (runtime *Runtime) respondWithPartialProgress( return nil, false, err } + bundle.ToolEvents = mergeNestedToolEvents(bundle.ToolEvents, progress.completedNestedTools) + return bundle, cached, nil } +type nestedToolEventMerger struct { + childrenByParent map[string][]*ToolEvent + nestedCallIDs map[string]struct{} + outerCallIDs map[string]struct{} + appended map[*ToolEvent]bool + visiting map[*ToolEvent]bool + merged []ToolEvent +} + +func mergeNestedToolEvents(outer, nested []ToolEvent) []ToolEvent { + if len(nested) == 0 { + return outer + } + + merger := newNestedToolEventMerger(outer, nested) + for index := range outer { + merger.appendChildren(outer[index].CallID) + merger.merged = append(merger.merged, outer[index]) + } + + for index := range nested { + if merger.isRootOrOrphan(&nested[index]) { + merger.appendNested(&nested[index]) + } + } + + for index := range nested { + merger.appendNested(&nested[index]) + } + + return merger.merged +} + +func newNestedToolEventMerger(outer, nested []ToolEvent) *nestedToolEventMerger { + merger := &nestedToolEventMerger{ + childrenByParent: make(map[string][]*ToolEvent), + nestedCallIDs: make(map[string]struct{}, len(nested)), + outerCallIDs: make(map[string]struct{}, len(outer)), + appended: make(map[*ToolEvent]bool, len(nested)), + visiting: make(map[*ToolEvent]bool, len(nested)), + merged: make([]ToolEvent, 0, len(outer)+len(nested)), + } + + for index := range outer { + merger.addOuter(&outer[index]) + } + + for index := range nested { + merger.addNested(&nested[index]) + } + + for parentID := range merger.childrenByParent { + slices.SortStableFunc(merger.childrenByParent[parentID], compareToolEventSequence) + } + + return merger +} + +func (merger *nestedToolEventMerger) addOuter(event *ToolEvent) { + if event.CallID != "" { + merger.outerCallIDs[event.CallID] = struct{}{} + } +} + +func (merger *nestedToolEventMerger) addNested(event *ToolEvent) { + if event.CallID != "" { + merger.nestedCallIDs[event.CallID] = struct{}{} + } + + if event.ParentCallID != "" { + merger.childrenByParent[event.ParentCallID] = append(merger.childrenByParent[event.ParentCallID], event) + } +} + +func compareToolEventSequence(left, right *ToolEvent) int { + switch { + case left.Sequence <= 0 && right.Sequence <= 0: + return 0 + case left.Sequence <= 0: + return 1 + case right.Sequence <= 0: + return -1 + case left.Sequence < right.Sequence: + return -1 + case left.Sequence > right.Sequence: + return 1 + default: + return 0 + } +} + +func (merger *nestedToolEventMerger) isRootOrOrphan(event *ToolEvent) bool { + if event.ParentCallID == "" { + return true + } + + _, parentIsNested := merger.nestedCallIDs[event.ParentCallID] + _, parentIsOuter := merger.outerCallIDs[event.ParentCallID] + + return !parentIsNested && !parentIsOuter +} + +func (merger *nestedToolEventMerger) appendChildren(parentCallID string) { + for _, child := range merger.childrenByParent[parentCallID] { + merger.appendNested(child) + } +} + +func (merger *nestedToolEventMerger) appendNested(event *ToolEvent) { + if merger.appended[event] || merger.visiting[event] { + return + } + + merger.visiting[event] = true + merger.appendChildren(event.CallID) + delete(merger.visiting, event) + merger.appended[event] = true + merger.merged = append(merger.merged, *event) +} + func newPartialPromptProgress(forward func(StreamEvent)) *partialPromptProgress { return &partialPromptProgress{ - forward: forward, - blocks: []partialPromptBlock{}, - fallbackBlocks: nil, - pendingTools: []pendingToolCall{}, + forward: forward, + blocks: []partialPromptBlock{}, + fallbackBlocks: nil, + pendingTools: []pendingToolCall{}, + completedNestedTools: []ToolEvent{}, } } @@ -154,6 +282,10 @@ func (progress *partialPromptProgress) record(streamEvent StreamEvent) { if streamEvent.ToolEvent != nil { progress.removePendingTool(streamEvent.ToolEvent) progress.append(transcript.RoleToolResult, formatToolEvent(streamEvent.ToolEvent)) + + if streamEvent.ToolEvent.ParentCallID != "" { + progress.completedNestedTools = append(progress.completedNestedTools, *streamEvent.ToolEvent) + } } case StreamEventToolStart: progress.trackPendingTool(streamEvent.ToolCallEvent, streamEvent.Text) @@ -192,6 +324,7 @@ func (progress *partialPromptProgress) reset() { progress.blocks = progress.blocks[:0] progress.pendingTools = progress.pendingTools[:0] + progress.completedNestedTools = progress.completedNestedTools[:0] } func (progress *partialPromptProgress) append(role transcript.Role, content string) { @@ -272,10 +405,10 @@ func (progress *partialPromptProgress) failureMessages(promptErr error) []partia }) } - for _, event := range toolEvents { + for index := range toolEvents { messages = append(messages, partialPromptMessage{ Role: database.RoleToolResult, - Content: formatToolEvent(&event), + Content: formatToolEvent(&toolEvents[index]), }) } @@ -292,10 +425,15 @@ func (progress *partialPromptProgress) trackPendingTool(call *ToolCallEvent, fal return } - pending := pendingToolCall{Name: fallbackName, ArgumentsJSON: ""} + pending := pendingToolCall{ + CallID: "", ParentCallID: "", Name: fallbackName, ArgumentsJSON: "", Sequence: 0, + } if call != nil { + pending.CallID = call.ID + pending.ParentCallID = call.ParentCallID pending.Name = call.Name pending.ArgumentsJSON = call.ArgumentsJSON + pending.Sequence = call.Sequence } if pending.Name == "" { @@ -314,7 +452,7 @@ func (progress *partialPromptProgress) removePendingTool(event *ToolEvent) { return } - index, found := progress.pendingToolIndex(event.Name, event.ArgumentsJSON) + index, found := progress.pendingToolIndex(event.CallID, event.Name, event.ArgumentsJSON) if !found { return } @@ -322,15 +460,40 @@ func (progress *partialPromptProgress) removePendingTool(event *ToolEvent) { progress.pendingTools = slices.Delete(progress.pendingTools, index, index+1) } -func (progress *partialPromptProgress) pendingToolIndex(name, argumentsJSON string) (int, bool) { - for index, pending := range progress.pendingTools { - if pending.Name == name && pending.ArgumentsJSON == argumentsJSON { +func (progress *partialPromptProgress) pendingToolIndex(callID, name, argumentsJSON string) (int, bool) { + if index, found := pendingToolIndexByCallID(progress.pendingTools, callID); found { + return index, true + } + + if index, found := pendingToolIndexByName(progress.pendingTools, callID, name, argumentsJSON); found { + return index, true + } + + return pendingToolIndexByName(progress.pendingTools, callID, name, "") +} + +func pendingToolIndexByCallID(pendingTools []pendingToolCall, callID string) (int, bool) { + if callID == "" { + return 0, false + } + + for index, pending := range pendingTools { + if pending.CallID == callID { return index, true } } - for index, pending := range progress.pendingTools { - if pending.Name == name { + return 0, false +} + +func pendingToolIndexByName(pendingTools []pendingToolCall, callID, name, argumentsJSON string) (int, bool) { + for index, pending := range pendingTools { + if callID != "" && pending.CallID != "" { + continue + } + + argumentsMatch := argumentsJSON == "" || pending.ArgumentsJSON == argumentsJSON + if pending.Name == name && argumentsMatch { return index, true } } @@ -348,6 +511,10 @@ func (progress *partialPromptProgress) syntheticToolFailureEvents(promptErr erro events := make([]ToolEvent, 0, len(progress.pendingTools)) for _, pending := range progress.pendingTools { events = append(events, ToolEvent{ + CallID: pending.CallID, + ParentCallID: pending.ParentCallID, + Sequence: pending.Sequence, + Name: pending.Name, ArgumentsJSON: pending.ArgumentsJSON, DetailsJSON: "", @@ -402,6 +569,9 @@ func formatToolEvent(toolEvent *ToolEvent) string { } return transcript.FormatToolEventPersistence(&transcript.ToolEvent{ + CallID: toolEvent.CallID, + ParentCallID: toolEvent.ParentCallID, + Sequence: toolEvent.Sequence, Name: toolEvent.Name, ArgumentsJSON: toolEvent.ArgumentsJSON, DetailsJSON: toolEvent.DetailsJSON, diff --git a/internal/assistant/runtime_persist_internal_test.go b/internal/assistant/runtime_persist_internal_test.go index d091fdb8..97776c31 100644 --- a/internal/assistant/runtime_persist_internal_test.go +++ b/internal/assistant/runtime_persist_internal_test.go @@ -5,16 +5,26 @@ import ( "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "github.com/omarluq/librecode/internal/tool" ) -const runtimePersistTestArgsJSON = `{"path":"a.go"}` +const ( + runtimePersistTestArgsJSON = `{"path":"a.go"}` + runtimePersistToolCallID = "call-1" + runtimePersistParentID = "outer" + runtimePersistChildOneID = "outer/1" +) func TestFormatToolEventIncludesErrorMarker(t *testing.T) { t.Parallel() formatted := formatToolEvent(&ToolEvent{ + CallID: "", + ParentCallID: "", + Sequence: 0, + Name: jsonBashToolName, ArgumentsJSON: `{"command":"false"}`, DetailsJSON: `{"exit_code":1}`, @@ -52,6 +62,118 @@ func TestPartialPromptProgressReplacesPendingToolWithResult(t *testing.T) { assert.Empty(t, progress.syntheticToolFailureEvents(context.Canceled)) } +func TestPartialPromptProgressFallsBackWhenResultHasGeneratedCallID(t *testing.T) { + t.Parallel() + + progress := newPartialPromptProgress(nil) + call := runtimePersistTestToolCall() + call.ID = "" + progress.trackPendingTool(call, "") + + result := runtimePersistTestToolResult() + result.CallID = "generated-call-id" + progress.removePendingTool(result) + + assert.Empty(t, progress.pendingTools) +} + +func TestPartialPromptProgressCollectsNestedToolsAndMatchesByCallID(t *testing.T) { + t.Parallel() + + progress := newPartialPromptProgress(nil) + first := runtimePersistTestToolCall() + first.ParentCallID = runtimePersistParentID + first.Sequence = 1 + second := runtimePersistTestToolCall() + second.ID = "call-2" + second.ParentCallID = runtimePersistParentID + second.Sequence = 2 + + progress.trackPendingTool(first, "") + progress.trackPendingTool(second, "") + + result := runtimePersistTestToolResult() + result.CallID = second.ID + result.ParentCallID = second.ParentCallID + result.Sequence = second.Sequence + progress.record(StreamEvent{ + ToolCallEvent: nil, ToolEvent: result, Usage: nil, Kind: StreamEventToolResult, Text: "", + }) + + require.Len(t, progress.pendingTools, 1) + assert.Equal(t, first.ID, progress.pendingTools[0].CallID) + require.Len(t, progress.completedNestedTools, 1) + assert.Equal(t, second.ID, progress.completedNestedTools[0].CallID) +} + +func TestMergeNestedToolEventsPreservesParentExecutionOrder(t *testing.T) { + t.Parallel() + + outer := []ToolEvent{ + runtimePersistIdentityEvent("outer-1", "", "execute", 0), + runtimePersistIdentityEvent("outer-2", "", "execute", 0), + } + nested := []ToolEvent{ + runtimePersistIdentityEvent("outer-1/1", "outer-1", "read", 1), + runtimePersistIdentityEvent("outer-2/1", "outer-2", "grep", 1), + runtimePersistIdentityEvent("outer-1/2", "outer-1", "find", 2), + } + + merged := mergeNestedToolEvents(outer, nested) + + require.Len(t, merged, 5) + assert.Equal(t, []string{"outer-1/1", "outer-1/2", "outer-1", "outer-2/1", "outer-2"}, []string{ + merged[0].CallID, merged[1].CallID, merged[2].CallID, merged[3].CallID, merged[4].CallID, + }) +} + +func TestMergeNestedToolEventsOrdersSiblingsBySequence(t *testing.T) { + t.Parallel() + + outer := []ToolEvent{runtimePersistIdentityEvent(runtimePersistParentID, "", "execute", 0)} + nested := []ToolEvent{ + runtimePersistIdentityEvent("outer/3", runtimePersistParentID, "write", 3), + runtimePersistIdentityEvent(runtimePersistChildOneID, runtimePersistParentID, "read", 1), + runtimePersistIdentityEvent("outer/unknown", runtimePersistParentID, "grep", 0), + runtimePersistIdentityEvent("outer/2", runtimePersistParentID, "find", 2), + } + + merged := mergeNestedToolEvents(outer, nested) + + wantCallIDs := []string{ + runtimePersistChildOneID, "outer/2", "outer/3", "outer/unknown", runtimePersistParentID, + } + assert.Equal(t, wantCallIDs, toolEventCallIDs(merged)) +} + +func TestMergeNestedToolEventsPreservesNestedHierarchy(t *testing.T) { + t.Parallel() + + outer := []ToolEvent{runtimePersistIdentityEvent(runtimePersistParentID, "", "execute", 0)} + nested := []ToolEvent{ + runtimePersistIdentityEvent("child", runtimePersistParentID, "execute", 1), + runtimePersistIdentityEvent("grandchild", "child", "read", 1), + } + + merged := mergeNestedToolEvents(outer, nested) + + assert.Equal(t, []string{"grandchild", "child", runtimePersistParentID}, toolEventCallIDs(merged)) +} + +func TestMergeNestedToolEventsKeepsOrphansInStreamOrder(t *testing.T) { + t.Parallel() + + nested := []ToolEvent{ + runtimePersistIdentityEvent("orphan-2", "missing", "grep", 2), + runtimePersistIdentityEvent("root", "", "read", 0), + runtimePersistIdentityEvent("orphan-1", "missing", "find", 1), + } + + merged := mergeNestedToolEvents(nil, nested) + + assert.Equal(t, []string{"orphan-2", "root", "orphan-1"}, toolEventCallIDs(merged)) +} + func TestPartialPromptProgressResetClearsPendingTools(t *testing.T) { t.Parallel() @@ -69,10 +191,29 @@ func TestPartialPromptProgressResetClearsPendingTools(t *testing.T) { assert.Empty(t, progress.syntheticToolFailureEvents(context.Canceled)) } +func toolEventCallIDs(events []ToolEvent) []string { + callIDs := make([]string, 0, len(events)) + for index := range events { + callIDs = append(callIDs, events[index].CallID) + } + + return callIDs +} + +func runtimePersistIdentityEvent(callID, parentCallID, name string, sequence int) ToolEvent { + return ToolEvent{ + CallID: callID, ParentCallID: parentCallID, Sequence: sequence, + Name: name, ArgumentsJSON: "", DetailsJSON: "", Result: "", Error: "", IsError: false, + } +} + func runtimePersistTestToolCall() *ToolCallEvent { return &ToolCallEvent{ + ParentCallID: "", + Sequence: 0, + ArgumentsJSON: runtimePersistTestArgsJSON, - ID: "call-1", + ID: runtimePersistToolCallID, Name: jsonReadToolName, Arguments: tool.EmptyArguments(), } @@ -80,6 +221,10 @@ func runtimePersistTestToolCall() *ToolCallEvent { func runtimePersistTestToolResult() *ToolEvent { return &ToolEvent{ + CallID: "", + ParentCallID: "", + Sequence: 0, + Name: jsonReadToolName, ArgumentsJSON: runtimePersistTestArgsJSON, DetailsJSON: "", diff --git a/internal/assistant/runtime_slash.go b/internal/assistant/runtime_slash.go index 1a3ab120..3486f0fb 100644 --- a/internal/assistant/runtime_slash.go +++ b/internal/assistant/runtime_slash.go @@ -109,22 +109,45 @@ func (runtime *Runtime) loadSkillWithReadTool( payload, err := skillReadArgumentsJSON(skill.FilePath, limit) if err != nil { - return "", ToolEvent{}, assistantError(err, "build skill read arguments") + return "", ToolEvent{ + CallID: "", + ParentCallID: "", + Name: "", + ArgumentsJSON: "", + DetailsJSON: "", + Result: "", + Error: "", + Sequence: 0, + IsError: false, + }, assistantError(err, "build skill read arguments") } arguments, err := tool.ArgumentsFromRaw(payload) if err != nil { - return "", ToolEvent{}, assistantError(err, "build skill read arguments") + return "", ToolEvent{ + CallID: "", + ParentCallID: "", + Name: "", + ArgumentsJSON: "", + DetailsJSON: "", + Result: "", + Error: "", + Sequence: 0, + IsError: false, + }, assistantError(err, "build skill read arguments") } result, err := registry.Execute(ctx, string(tool.NameRead), arguments) toolEvent = ToolEvent{ + CallID: "", + ParentCallID: "", Name: "load skill: " + skill.Name, ArgumentsJSON: string(payload), DetailsJSON: "", Result: "", Error: "", + Sequence: 0, IsError: false, } if err != nil { diff --git a/internal/assistant/runtime_test.go b/internal/assistant/runtime_test.go index ef571de1..6fca336b 100644 --- a/internal/assistant/runtime_test.go +++ b/internal/assistant/runtime_test.go @@ -12,6 +12,7 @@ import ( "testing" "time" + "github.com/samber/oops" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" _ "modernc.org/sqlite" // register sqlite driver for assistant runtime tests @@ -73,6 +74,38 @@ func newRuntimeOutsideTempDir(t *testing.T) string { return dir } +func TestRuntime_PromptPersistsSuccessfulNestedToolsInExecutionGroups(t *testing.T) { + t.Parallel() + + cwd := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(cwd, "one.txt"), []byte("one"), 0o600)) + require.NoError(t, os.WriteFile(filepath.Join(cwd, "two.txt"), []byte("two"), 0o600)) + + runtime, repository := newTestRuntimeWithClient(t, nestedExecuteCompleter{}) + request := newRuntimePromptRequest(cwd, "nested persistence", "") + response, err := runtime.Prompt(context.Background(), request) + require.NoError(t, err) + + require.Len(t, response.ToolEvents, 3) + assert.Equal(t, []string{"execute-call/1", "execute-call/2", "execute-call"}, []string{ + response.ToolEvents[0].CallID, + response.ToolEvents[1].CallID, + response.ToolEvents[2].CallID, + }) + + messages := requireRuntimeMessages(t, repository, response.SessionID, 5) + assert.Equal(t, database.RoleToolResult, messages[1].Role) + assert.Contains(t, messages[1].Content, "call_id: execute-call/1") + assert.Contains(t, messages[1].Content, "parent_call_id: execute-call") + assert.Equal(t, database.RoleToolResult, messages[2].Role) + assert.Contains(t, messages[2].Content, "call_id: execute-call/2") + assert.Contains(t, messages[2].Content, "parent_call_id: execute-call") + assert.Equal(t, database.RoleToolResult, messages[3].Role) + assert.Contains(t, messages[3].Content, "call_id: execute-call") + assert.NotContains(t, messages[3].Content, "parent_call_id:") + assert.Equal(t, database.RoleAssistant, messages[4].Role) +} + func TestRuntime_PromptPersistsConversation(t *testing.T) { t.Parallel() @@ -861,6 +894,9 @@ func emitPendingToolStart(request *assistant.CompletionRequest) { emitPartialText(request) request.OnEvent(assistant.StreamEvent{ ToolCallEvent: &assistant.ToolCallEvent{ + ParentCallID: "", + Sequence: 0, + ArgumentsJSON: testToolArgsJSON, ID: testToolCallID, Name: testToolName, @@ -882,6 +918,10 @@ func emitPartialToolResult(request *assistant.CompletionRequest) { request.OnEvent(assistant.StreamEvent{ ToolCallEvent: nil, ToolEvent: &assistant.ToolEvent{ + CallID: "", + ParentCallID: "", + Sequence: 0, + Name: testToolName, ArgumentsJSON: testToolArgsJSON, DetailsJSON: "", @@ -964,6 +1004,33 @@ func assertRuntimeMessage( } } +type nestedExecuteCompleter struct{} + +func (nestedExecuteCompleter) Complete( + ctx context.Context, + request *assistant.CompletionRequest, +) (*assistant.CompletionResult, error) { + source := `import "tools" +tools.Call("read", map[string]interface{}{"path": "one.txt"}) +tools.Call("read", map[string]interface{}{"path": "two.txt"})` + arguments := testutil.ToolArguments(map[string]any{"source": source}) + + events, err := request.ExecuteTools(ctx, []assistant.ToolCall{{ + Metadata: nil, ArgumentsJSON: arguments.String(), ID: "execute-call", Name: "execute", Arguments: arguments, + }}, request.OnEvent) + if err != nil { + return nil, oops.In("assistant_test").Code("execute_nested_tools").Wrapf(err, "execute nested tools") + } + + return &assistant.CompletionResult{ + FinishReason: llm.FinishReasonStop, + Text: "done", + Thinking: nil, + ToolEvents: events, + Usage: model.EmptyTokenUsage(), + }, nil +} + type testCompleter struct{} func (testCompleter) Complete( diff --git a/internal/assistant/tool_arguments_internal_test.go b/internal/assistant/tool_arguments_internal_test.go index a301f838..d19b29df 100644 --- a/internal/assistant/tool_arguments_internal_test.go +++ b/internal/assistant/tool_arguments_internal_test.go @@ -15,6 +15,9 @@ func TestApplyToolCallMutationAppliesEmptyArguments(t *testing.T) { t.Parallel() call := &ToolCallEvent{ + ParentCallID: "", + Sequence: 0, + Arguments: testutil.ToolArguments(map[string]any{"path": "old.txt"}), ArgumentsJSON: `{"path":"old.txt"}`, ID: "", diff --git a/internal/assistant/tool_executor.go b/internal/assistant/tool_executor.go index 1d966035..77ecb703 100644 --- a/internal/assistant/tool_executor.go +++ b/internal/assistant/tool_executor.go @@ -3,8 +3,11 @@ package assistant import ( "context" "encoding/json" + "fmt" "strings" + "sync" + "github.com/gofrs/uuid/v5" "github.com/samber/oops" "github.com/omarluq/librecode/internal/llm" @@ -32,6 +35,34 @@ func (runtime *Runtime) executeProviderToolCalls( } } +type toolInvocationScope struct { + onEvent func(StreamEvent) + parentCallID string + nextSequence int + mu sync.Mutex +} + +func (scope *toolInvocationScope) nextCall(name string) ToolCallEvent { + scope.mu.Lock() + defer scope.mu.Unlock() + + scope.nextSequence++ + sequence := scope.nextSequence + + callID := fmt.Sprintf("%s/%d", scope.parentCallID, sequence) + + return ToolCallEvent{ + ArgumentsJSON: "", + ID: callID, + ParentCallID: scope.parentCallID, + Name: name, + Arguments: tool.EmptyArguments(), + Sequence: sequence, + } +} + +type toolInvocationContextKey struct{} + func (runtime *Runtime) executeProviderToolCall( ctx context.Context, registry *tool.Registry, @@ -39,28 +70,51 @@ func (runtime *Runtime) executeProviderToolCall( onEvent func(StreamEvent), ) ToolEvent { callEvent := ToolCallEvent{ - Arguments: tool.EmptyArguments(), ArgumentsJSON: "", ID: "", + ParentCallID: "", Name: "", + Arguments: tool.EmptyArguments(), + Sequence: 0, } if call != nil { callEvent.Arguments = call.Arguments callEvent.ID = call.ID + callEvent.ParentCallID = stringFromOptions(call.Metadata, toolParentCallIDMetadataKey) callEvent.Name = call.Name callEvent.ArgumentsJSON = call.ArgumentsJSON + callEvent.Sequence = sequenceFromOptions(call.Metadata) + } + + if callEvent.ID == "" { + callEvent.ID = uuid.Must(uuid.NewV7()).String() } - if err := runtime.dispatchToolCallLifecycle(ctx, &callEvent); err != nil { + event, _ := runtime.invokeToolResult(ctx, registry, &callEvent, onEvent) + + return event +} + +func (runtime *Runtime) invokeToolResult( + ctx context.Context, + registry *tool.Registry, + callEvent *ToolCallEvent, + onEvent func(StreamEvent), +) (ToolEvent, tool.Result) { + if err := runtime.dispatchToolCallLifecycle(ctx, callEvent); err != nil { event := toolLifecycleErrorEvent(callEvent, err) emitProviderToolResult(onEvent, &event) - return event + return event, tool.Result{Content: nil, Details: nil} } - emitProviderToolStart(onEvent, &callEvent) + emitProviderToolStart(onEvent, callEvent) - result, err := registry.Execute(ctx, callEvent.Name, callEvent.Arguments) + scope := &toolInvocationScope{ + onEvent: onEvent, parentCallID: callEvent.ID, nextSequence: 0, mu: sync.Mutex{}, + } + nestedCtx := context.WithValue(ctx, toolInvocationContextKey{}, scope) + result, err := registry.Execute(nestedCtx, callEvent.Name, callEvent.Arguments) event := toolEventFromResult(callEvent, result, err) if lifecycleErr := runtime.dispatchToolResultLifecycle(ctx, &event); lifecycleErr != nil && runtime.logger != nil { @@ -68,22 +122,67 @@ func (runtime *Runtime) executeProviderToolCall( } emitProviderToolResult(onEvent, &event) + result = canonicalToolResult(result, &event) - return event + return event, result } -func toolLifecycleErrorEvent(call ToolCallEvent, err error) ToolEvent { +func (runtime *Runtime) invokeNestedTool( + ctx context.Context, + registry *tool.Registry, + name string, + arguments tool.Arguments, + argumentsJSON string, +) (tool.Result, ToolEvent) { + scope := toolInvocationScopeFromContext(ctx) + call := ToolCallEvent{ + ArgumentsJSON: argumentsJSON, + ID: "", + ParentCallID: "", + Name: name, + Arguments: arguments, + Sequence: 0, + } + + var onEvent func(StreamEvent) + + if scope != nil { + identity := scope.nextCall(name) + call.ID = identity.ID + call.ParentCallID = identity.ParentCallID + call.Sequence = identity.Sequence + onEvent = scope.onEvent + } + + event, result := runtime.invokeToolResult(ctx, registry, &call, onEvent) + + return result, event +} + +func toolInvocationScopeFromContext(ctx context.Context) *toolInvocationScope { + scope, ok := ctx.Value(toolInvocationContextKey{}).(*toolInvocationScope) + if !ok { + return nil + } + + return scope +} + +func toolLifecycleErrorEvent(call *ToolCallEvent, err error) ToolEvent { return ToolEvent{ + CallID: call.ID, + ParentCallID: call.ParentCallID, Name: call.Name, ArgumentsJSON: call.ArgumentsJSON, DetailsJSON: "", Result: err.Error(), Error: err.Error(), + Sequence: call.Sequence, IsError: true, } } -func toolEventFromResult(call ToolCallEvent, result tool.Result, err error) ToolEvent { +func toolEventFromResult(call *ToolCallEvent, result tool.Result, err error) ToolEvent { resultText := result.Text() detailsJSON := encodeToolDetails(result.Details) @@ -96,11 +195,14 @@ func toolEventFromResult(call ToolCallEvent, result tool.Result, err error) Tool } event := ToolEvent{ + CallID: call.ID, + ParentCallID: call.ParentCallID, Name: call.Name, ArgumentsJSON: call.ArgumentsJSON, DetailsJSON: detailsJSON, Result: resultText, Error: "", + Sequence: call.Sequence, IsError: false, } if err != nil { @@ -111,6 +213,31 @@ func toolEventFromResult(call ToolCallEvent, result tool.Result, err error) Tool return event } +func canonicalToolResult(result tool.Result, event *ToolEvent) tool.Result { + if event == nil { + return result + } + + if event.Result != result.Text() { + result.Content = []tool.ContentBlock{{Type: tool.ContentTypeText, Text: event.Result, Data: "", MIMEType: ""}} + } + + if strings.TrimSpace(event.DetailsJSON) == "" { + result.Details = nil + + return result + } + + result.Details = nil + + var details map[string]any + if err := json.Unmarshal([]byte(event.DetailsJSON), &details); err == nil { + result.Details = details + } + + return result +} + func encodeToolDetails(details map[string]any) string { if len(details) == 0 { return "" @@ -161,7 +288,7 @@ func llmToolResultFromToolEvent(event *ToolEvent) llm.ToolResult { return llm.ToolResult{ Metadata: toolResultMetadataFromToolEvent(event), - ToolCallID: "", + ToolCallID: event.CallID, ArgumentsJSON: event.ArgumentsJSON, Name: event.Name, Error: event.Error, @@ -171,9 +298,20 @@ func llmToolResultFromToolEvent(event *ToolEvent) llm.ToolResult { } func toolResultMetadataFromToolEvent(event *ToolEvent) map[string]any { - if event == nil || strings.TrimSpace(event.DetailsJSON) == "" { + if event == nil { return nil } - return map[string]any{"details_json": event.DetailsJSON} + metadata := toolIdentityMetadata(event.ParentCallID, event.Sequence) + if strings.TrimSpace(event.DetailsJSON) == "" { + return metadata + } + + if metadata == nil { + metadata = make(map[string]any, 1) + } + + metadata["details_json"] = event.DetailsJSON + + return metadata } diff --git a/internal/assistant/tool_executor_internal_test.go b/internal/assistant/tool_executor_internal_test.go index 7f2107d7..91d0af1b 100644 --- a/internal/assistant/tool_executor_internal_test.go +++ b/internal/assistant/tool_executor_internal_test.go @@ -164,6 +164,35 @@ func TestExecuteProviderToolCallPreservesResultOnLifecycleError(t *testing.T) { assert.Equal(t, event.Result, events[1].ToolEvent.Result) } +func TestCanonicalToolResultUsesLifecycleMutation(t *testing.T) { + t.Parallel() + + tests := []struct { + wantDetails map[string]any + name string + detailsJSON string + }{ + {name: "valid details", detailsJSON: `{"redacted":true}`, wantDetails: map[string]any{"redacted": true}}, + {name: "malformed details fail closed", detailsJSON: `{"redacted":`, wantDetails: nil}, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + result := tool.TextResult("original", map[string]any{"old": true}) + event := ToolEvent{ + CallID: executeCallMethod, ParentCallID: "", Name: jsonReadToolName, ArgumentsJSON: "{}", + DetailsJSON: testCase.detailsJSON, Result: "redacted", Error: "", Sequence: 0, IsError: false, + } + + canonical := canonicalToolResult(result, &event) + assert.Equal(t, "redacted", canonical.Text()) + assert.Equal(t, testCase.wantDetails, canonical.Details) + }) + } +} + func TestEncodeToolDetails(t *testing.T) { t.Parallel() @@ -176,7 +205,10 @@ func TestToolEventFromResultFormatsEmptyOutput(t *testing.T) { t.Parallel() event := toolEventFromResult( - ToolCallEvent{ + &ToolCallEvent{ + ParentCallID: "", + Sequence: 0, + Arguments: tool.EmptyArguments(), ID: "", Name: jsonReadToolName, @@ -198,6 +230,10 @@ func TestLLMToolResultFromToolEvent(t *testing.T) { assert.False(t, empty.IsError) result := llmToolResultFromToolEvent(&ToolEvent{ + CallID: "", + ParentCallID: "", + Sequence: 0, + Name: jsonReadToolName, ArgumentsJSON: toolExecutorReadArgs, DetailsJSON: `{"diff":"+added"}`, diff --git a/internal/assistant/tool_lifecycle_test.go b/internal/assistant/tool_lifecycle_test.go index 122d295c..3709e26c 100644 --- a/internal/assistant/tool_lifecycle_test.go +++ b/internal/assistant/tool_lifecycle_test.go @@ -12,7 +12,11 @@ import ( "github.com/omarluq/librecode/internal/tool" ) -const testToolLifecycleError = "boom" +const ( + testToolLifecycleError = "boom" + testToolLifecycleParentID = "parent-call" + testToolLifecycleSequence = 7 +) func TestRuntime_ToolCallLifecycleAppliesArgumentMutation(t *testing.T) { t.Parallel() @@ -30,6 +34,9 @@ func TestRuntime_ToolCallLifecycleAppliesArgumentMutation(t *testing.T) { lua: ` local lc = require("librecode") lc.on("tool_call", function(event) + assert(event.payload.call_id == "call-1") + assert(event.payload.parent_call_id == "parent-call") + assert(event.payload.sequence == 7) return { tool_call = { arguments = { @@ -55,6 +62,9 @@ end) runtime, _, manager := newTestRuntimeWithManager(t, testCompleter{}) loadRuntimeExtension(t, manager, testCase.lua) call := assistant.ToolCallEvent{ + ParentCallID: testToolLifecycleParentID, + Sequence: testToolLifecycleSequence, + Arguments: testCase.initialArguments, ID: testToolCallID, Name: testToolName, @@ -70,6 +80,9 @@ end) testutil.ToolArgumentFields(call.Arguments), ) assert.JSONEq(t, testCase.expectedArgumentsJSON, call.ArgumentsJSON) + assert.Equal(t, testToolCallID, call.ID) + assert.Equal(t, testToolLifecycleParentID, call.ParentCallID) + assert.Equal(t, testToolLifecycleSequence, call.Sequence) }) } } @@ -88,6 +101,10 @@ func TestRuntime_ToolResultLifecycleAppliesResultMutation(t *testing.T) { { name: "redacts result and clears error", initialEvent: &assistant.ToolEvent{ + CallID: testToolCallID, + ParentCallID: testToolLifecycleParentID, + Sequence: testToolLifecycleSequence, + Name: testToolName, ArgumentsJSON: testToolArgsJSON, DetailsJSON: "{}", @@ -98,6 +115,9 @@ func TestRuntime_ToolResultLifecycleAppliesResultMutation(t *testing.T) { lua: ` local lc = require("librecode") lc.on("tool_result", function(event) + assert(event.payload.call_id == "call-1") + assert(event.payload.parent_call_id == "parent-call") + assert(event.payload.sequence == 7) return { tool_result = { result = "redacted", @@ -126,6 +146,9 @@ end) assert.Equal(t, testCase.expectedResult, testCase.initialEvent.Result) assert.JSONEq(t, testCase.expectedDetailsJSON, testCase.initialEvent.DetailsJSON) assert.Equal(t, testCase.expectedError, testCase.initialEvent.Error) + assert.Equal(t, testToolCallID, testCase.initialEvent.CallID) + assert.Equal(t, testToolLifecycleParentID, testCase.initialEvent.ParentCallID) + assert.Equal(t, testToolLifecycleSequence, testCase.initialEvent.Sequence) }) } } @@ -145,12 +168,19 @@ end) `) call := assistant.ToolCallEvent{ + ParentCallID: "", + Sequence: 0, + Arguments: testutil.ToolArguments(map[string]any{testToolPathKey: testToolPath}), ID: testToolCallID, Name: testToolName, ArgumentsJSON: testToolArgsJSON, } event := &assistant.ToolEvent{ + CallID: "", + ParentCallID: "", + Sequence: 0, + Name: testToolName, ArgumentsJSON: testToolArgsJSON, DetailsJSON: "{}", @@ -184,6 +214,10 @@ end) `) event := &assistant.ToolEvent{ + CallID: "", + ParentCallID: "", + Sequence: 0, + Name: testToolName, ArgumentsJSON: testToolArgsJSON, DetailsJSON: "", diff --git a/internal/assistant/tool_registry.go b/internal/assistant/tool_registry.go index d05ff188..8992390b 100644 --- a/internal/assistant/tool_registry.go +++ b/internal/assistant/tool_registry.go @@ -1,6 +1,7 @@ package assistant import ( + "context" "reflect" "github.com/samber/oops" @@ -40,7 +41,11 @@ func newToolRegistry(cwd string, provider toolProvider) (*tool.Registry, error) return registry, nil } -func (runtime *Runtime) promptToolRegistry(cwd, sessionID string) (*tool.Registry, error) { +func (runtime *Runtime) promptToolRegistry( + ctx context.Context, + cwd string, + sessionID string, +) (*tool.Registry, error) { if runtime.profile.Kind != ExecutionTopLevel { return runtime.profileToolRegistry(cwd) } @@ -55,32 +60,76 @@ func (runtime *Runtime) promptToolRegistry(cwd, sessionID string) (*tool.Registr return nil, err } - if runtime.agents != nil && runtime.agentTasks != nil && runtime.agents.Len() > 0 { - for _, name := range []tool.Name{ - agentStartToolName, - agentStatusToolName, - agentWaitToolName, - agentCancelToolName, - agentListToolName, - } { - executor := &agentToolExecutor{ - controller: runtime.agentTasks, - sessions: runtime.sessions, - catalog: runtime.agents, - name: name, - parentSessionID: sessionID, - cwd: cwd, - definition: nil, - } - if err := registry.Register(executor); err != nil { - return nil, oops.In("assistant").Code("register_agent_tool").Wrapf(err, "register agent tool") - } - } + if err := runtime.registerAgentTools(registry, sessionID, cwd); err != nil { + return nil, err + } + + if err := runtime.registerWorkflowTool(registry, sessionID); err != nil { + return nil, err + } + + if err := runtime.registerExecuteTool(ctx, registry); err != nil { + return nil, err } return registry, nil } +func (runtime *Runtime) registerAgentTools(registry *tool.Registry, sessionID, cwd string) error { + if runtime.agents == nil || runtime.agentTasks == nil || runtime.agents.Len() == 0 { + return nil + } + + for _, name := range []tool.Name{ + agentStartToolName, agentStatusToolName, agentWaitToolName, agentCancelToolName, agentListToolName, + } { + executor := &agentToolExecutor{ + controller: runtime.agentTasks, sessions: runtime.sessions, catalog: runtime.agents, name: name, + parentSessionID: sessionID, cwd: cwd, definition: nil, + } + if err := registry.Register(executor); err != nil { + return oops.In("assistant").Code("register_agent_tool").Wrapf(err, "register agent tool") + } + } + + return nil +} + +func (runtime *Runtime) registerWorkflowTool(registry *tool.Registry, sessionID string) error { + if runtime.workflowSubmitter == nil { + return nil + } + + executor := &workflowToolExecutor{ + submitter: runtime.workflowSubmitter, ownerSessionID: sessionID, + } + if err := registry.Register(executor); err != nil { + return oops.In("assistant").Code("register_workflow_tool").Wrapf(err, "register workflow tool") + } + + return nil +} + +func (runtime *Runtime) registerExecuteTool(ctx context.Context, registry *tool.Registry) error { + if toolStrategyFromContext(ctx) == ToolStrategyDirect { + return nil + } + + invoke := func( + ctx context.Context, + name string, + arguments tool.Arguments, + argumentsJSON string, + ) (tool.Result, ToolEvent) { + return runtime.invokeNestedTool(ctx, registry, name, arguments, argumentsJSON) + } + if err := registry.Register(newExecuteTool(registry, invoke)); err != nil { + return oops.In("assistant").Code("register_execute_tool").Wrapf(err, "register execute tool") + } + + return nil +} + func (runtime *Runtime) profileToolRegistry(cwd string) (*tool.Registry, error) { registry, err := tool.NewRegistryWithTools(cwd, runtime.profile.Tools) if err != nil { diff --git a/internal/assistant/tool_schema_cache_internal_test.go b/internal/assistant/tool_schema_cache_internal_test.go index 029346c6..e176675f 100644 --- a/internal/assistant/tool_schema_cache_internal_test.go +++ b/internal/assistant/tool_schema_cache_internal_test.go @@ -160,18 +160,19 @@ func TestRuntimeEstimateToolSchemaTokensCreatesRegistryWhenNil(t *testing.T) { func newTestRuntimeWithSchemaCache() *Runtime { return &Runtime{ - cfg: nil, - sessions: nil, - extensions: nil, - cache: nil, - models: nil, - client: nil, - logger: nil, - skillsCache: nil, - toolSchemaCache: newToolSchemaCache(), - agents: nil, - agentTasks: nil, - profile: topLevelExecutionProfile(), + cfg: nil, + sessions: nil, + extensions: nil, + cache: nil, + models: nil, + client: nil, + logger: nil, + skillsCache: nil, + toolSchemaCache: newToolSchemaCache(), + agents: nil, + agentTasks: nil, + workflowSubmitter: nil, + profile: topLevelExecutionProfile(), } } diff --git a/internal/assistant/workflow_controller.go b/internal/assistant/workflow_controller.go new file mode 100644 index 00000000..7bfef2df --- /dev/null +++ b/internal/assistant/workflow_controller.go @@ -0,0 +1,127 @@ +package assistant + +import ( + "context" + + "github.com/samber/oops" + + "github.com/omarluq/librecode/internal/database" + "github.com/omarluq/librecode/internal/workflow" +) + +// WorkflowController adapts the reusable agent submitter and task controller to +// the workflow runtime without weakening child-session ownership or policy +// snapshot semantics. +type WorkflowController struct { + submitter *AgentSubmitter + tasks AgentTaskController + sessions *database.SessionRepository +} + +// NewWorkflowController creates the production workflow agent adapter. +func NewWorkflowController( + submitter *AgentSubmitter, + tasks AgentTaskController, + sessions *database.SessionRepository, +) (*WorkflowController, error) { + if submitter == nil || tasks == nil || sessions == nil { + return nil, oops.In("assistant").Code("invalid_workflow_controller_dependencies"). + Errorf("agent submitter, task controller, and sessions are required") + } + + return &WorkflowController{submitter: submitter, tasks: tasks, sessions: sessions}, nil +} + +// Submit resolves the owner's working directory and delegates child creation +// and policy capture to AgentSubmitter. +func (controller *WorkflowController) Submit( + ctx context.Context, + request *workflow.AgentRequest, +) (*database.AgentTaskEntity, error) { + if request == nil { + return nil, oops.In("assistant").Code("nil_workflow_agent_request").Errorf("workflow agent request is nil") + } + + owner, found, err := controller.sessions.GetSession(ctx, request.OwnerSessionID) + if err != nil { + return nil, oops.In("assistant").Code("load_workflow_owner").Wrapf(err, "load workflow owner session") + } + + if !found { + return nil, oops.In("assistant").Code("workflow_owner_not_found"). + Errorf("workflow owner session %q was not found", request.OwnerSessionID) + } + + return controller.submitter.SubmitAgent(ctx, &AgentSubmitRequest{ + ParentTaskID: request.ParentTaskID, + OwnerSessionID: request.OwnerSessionID, + CWD: owner.CWD, + AgentName: request.Options.AgentName, + Prompt: request.Prompt, + Model: request.Options.Model, + Provider: request.Options.Provider, + ConcurrencyKey: request.Options.ConcurrencyKey, + NodeKey: request.NodeKey, + InvocationIndex: request.InvocationIndex, + Depth: request.Options.Depth, + }) +} + +// Get returns an agent task by ID. +func (controller *WorkflowController) Get( + ctx context.Context, + taskID string, +) (*database.AgentTaskEntity, bool, error) { + task, found, err := controller.tasks.Get(ctx, taskID) + if err != nil { + return nil, false, oops.In("assistant").Code("get_workflow_agent_task").Wrapf(err, "get workflow agent task") + } + + return task, found, nil +} + +// List returns agent tasks owned by a workflow session. +func (controller *WorkflowController) List( + ctx context.Context, + ownerSessionID string, + limit int, +) ([]database.AgentTaskEntity, error) { + tasks, err := controller.tasks.List(ctx, ownerSessionID, limit) + if err != nil { + return nil, oops.In("assistant").Code("list_workflow_agent_tasks").Wrapf(err, "list workflow agent tasks") + } + + return tasks, nil +} + +// Await waits for an agent task to reach a terminal state. +func (controller *WorkflowController) Await( + ctx context.Context, + taskID string, +) (*database.AgentTaskEntity, error) { + task, err := controller.tasks.Await(ctx, taskID) + if err != nil { + return nil, oops.In("assistant").Code("await_workflow_agent_task").Wrapf(err, "await workflow agent task") + } + + return task, nil +} + +// Cancel requests cancellation of a workflow-owned agent task. +func (controller *WorkflowController) Cancel( + ctx context.Context, + ownerSessionID string, + taskID string, +) (*database.TaskEntity, bool, error) { + task, found, err := controller.tasks.Cancel(ctx, ownerSessionID, taskID) + if err != nil { + return nil, false, oops.In("assistant").Code("cancel_workflow_agent_task").Wrapf( + err, + "cancel workflow agent task", + ) + } + + return task, found, nil +} + +var _ workflow.Controller = (*WorkflowController)(nil) diff --git a/internal/assistant/workflow_controller_internal_test.go b/internal/assistant/workflow_controller_internal_test.go new file mode 100644 index 00000000..ddb1d079 --- /dev/null +++ b/internal/assistant/workflow_controller_internal_test.go @@ -0,0 +1,228 @@ +package assistant + +import ( + "context" + "errors" + "testing" + + "github.com/samber/oops" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/omarluq/librecode/internal/database" + "github.com/omarluq/librecode/internal/workflow" +) + +const workflowControllerOwnerID = "workflow-controller-owner" + +type workflowControllerTaskStub struct { + task *database.AgentTaskEntity + request *AgentTaskRequest + err error + listed []database.AgentTaskEntity + found bool +} + +func (stub *workflowControllerTaskStub) SubmitAgentTask( + _ context.Context, + request *AgentTaskRequest, +) (*database.AgentTaskEntity, error) { + stub.request = request + + return stub.task, stub.err +} + +func (stub *workflowControllerTaskStub) Get( + context.Context, + string, +) (*database.AgentTaskEntity, bool, error) { + return stub.task, stub.found, stub.err +} + +func (stub *workflowControllerTaskStub) List( + context.Context, + string, + int, +) ([]database.AgentTaskEntity, error) { + return stub.listed, stub.err +} + +func (stub *workflowControllerTaskStub) Cancel( + context.Context, + string, + string, +) (*database.TaskEntity, bool, error) { + if stub.task == nil { + return nil, stub.found, stub.err + } + + return &stub.task.Task, stub.found, stub.err +} + +func (stub *workflowControllerTaskStub) Await( + context.Context, + string, +) (*database.AgentTaskEntity, error) { + return stub.task, stub.err +} + +func (*workflowControllerTaskStub) SubscribeAgentTask( + string, +) (events <-chan database.TaskEventEntity, unsubscribe func()) { + eventChannel := make(chan database.TaskEventEntity) + close(eventChannel) + + return eventChannel, func() {} +} + +func TestNewWorkflowControllerValidatesDependencies(t *testing.T) { + t.Parallel() + + tasks := new(workflowControllerTaskStub) + sessions := agentToolSessions(t) + submitter, err := NewAgentSubmitter(tasks, isolatedAgentCatalog(t)) + require.NoError(t, err) + + for _, testCase := range []struct { + tasks AgentTaskController + submitter *AgentSubmitter + sessions *database.SessionRepository + name string + }{ + {name: "nil submitter", tasks: tasks, submitter: nil, sessions: sessions}, + {name: "nil tasks", tasks: nil, submitter: submitter, sessions: sessions}, + {name: "nil sessions", tasks: tasks, submitter: submitter, sessions: nil}, + } { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + _, controllerErr := NewWorkflowController(testCase.submitter, testCase.tasks, testCase.sessions) + assertOopsCode(t, controllerErr, "invalid_workflow_controller_dependencies") + }) + } + + controller, err := NewWorkflowController(submitter, tasks, sessions) + require.NoError(t, err) + assert.Same(t, submitter, controller.submitter) + assert.Same(t, tasks, controller.tasks) + assert.Same(t, sessions, controller.sessions) +} + +func TestWorkflowControllerSubmitResolvesOwnerAndPreservesRequestIdentity(t *testing.T) { + t.Parallel() + + tasks := &workflowControllerTaskStub{ + err: nil, + task: agentToolTask("workflow-agent-task", workflowControllerOwnerID, database.TaskQueued), + request: nil, + listed: nil, + found: true, + } + sessions := agentToolSessions(t) + owner, err := sessions.CreateSession(t.Context(), t.TempDir(), "owner", "") + require.NoError(t, err) + submitter, err := NewAgentSubmitter(tasks, isolatedAgentCatalog(t)) + require.NoError(t, err) + controller, err := NewWorkflowController(submitter, tasks, sessions) + require.NoError(t, err) + + request := &workflow.AgentRequest{ + ParentTaskID: "parent-task", OwnerSessionID: owner.ID, NodeKey: "node-1", Prompt: "review", + Options: workflow.AgentOptions{ + NodeKey: "option-node", AgentName: defaultWorkflowAgentName, Model: "model", Provider: "provider", + ConcurrencyKey: "concurrency", Depth: 2, + }, + InvocationIndex: 3, + } + + task, err := controller.Submit(t.Context(), request) + require.NoError(t, err) + assert.Same(t, tasks.task, task) + require.NotNil(t, tasks.request) + assert.Equal(t, request.ParentTaskID, tasks.request.ParentTaskID) + assert.Equal(t, request.OwnerSessionID, tasks.request.OwnerSessionID) + assert.Equal(t, request.NodeKey, tasks.request.NodeKey) + assert.Equal(t, request.InvocationIndex, tasks.request.InvocationIndex) + assert.Equal(t, owner.CWD, tasks.request.ChildSessionCWD) +} + +func TestWorkflowControllerSubmitErrors(t *testing.T) { + t.Parallel() + + tasks := new(workflowControllerTaskStub) + sessions, databaseConnection := agentToolSessionsWithDB(t) + submitter, err := NewAgentSubmitter(tasks, isolatedAgentCatalog(t)) + require.NoError(t, err) + controller, err := NewWorkflowController(submitter, tasks, sessions) + require.NoError(t, err) + + _, err = controller.Submit(t.Context(), nil) + assertOopsCode(t, err, "nil_workflow_agent_request") + _, err = controller.Submit(t.Context(), emptyWorkflowAgentRequest(runtimeSlashMissing)) + assertOopsCode(t, err, "workflow_owner_not_found") + require.NoError(t, databaseConnection.Close()) + _, err = controller.Submit(t.Context(), emptyWorkflowAgentRequest("owner")) + assertOopsCode(t, err, "get_session") + assert.ErrorContains(t, err, "load workflow owner session") +} + +func TestWorkflowControllerDelegatesTaskOperations(t *testing.T) { + t.Parallel() + + task := agentToolTask("task-1", workflowControllerOwnerID, database.TaskRunning) + listed := []database.AgentTaskEntity{*task} + tasks := &workflowControllerTaskStub{err: nil, task: task, request: nil, listed: listed, found: true} + controller := &WorkflowController{submitter: nil, tasks: tasks, sessions: nil} + + got, found, err := controller.Get(t.Context(), task.Task.ID) + require.NoError(t, err) + assert.True(t, found) + assert.Same(t, task, got) + gotList, err := controller.List(t.Context(), workflowControllerOwnerID, 1) + require.NoError(t, err) + assert.Equal(t, listed, gotList) + awaited, err := controller.Await(t.Context(), task.Task.ID) + require.NoError(t, err) + assert.Same(t, task, awaited) + canceled, found, err := controller.Cancel(t.Context(), workflowControllerOwnerID, task.Task.ID) + require.NoError(t, err) + assert.True(t, found) + assert.Same(t, &task.Task, canceled) +} + +func TestWorkflowControllerWrapsTaskOperationErrors(t *testing.T) { + t.Parallel() + + tasks := &workflowControllerTaskStub{ + err: errors.New("backend failed"), task: nil, request: nil, listed: nil, found: false, + } + controller := &WorkflowController{submitter: nil, tasks: tasks, sessions: nil} + + _, _, err := controller.Get(t.Context(), "task") + assertOopsCode(t, err, "get_workflow_agent_task") + _, err = controller.List(t.Context(), workflowControllerOwnerID, 1) + assertOopsCode(t, err, "list_workflow_agent_tasks") + _, err = controller.Await(t.Context(), "task") + assertOopsCode(t, err, "await_workflow_agent_task") + _, _, err = controller.Cancel(t.Context(), workflowControllerOwnerID, "task") + assertOopsCode(t, err, "cancel_workflow_agent_task") +} + +func emptyWorkflowAgentRequest(owner string) *workflow.AgentRequest { + return &workflow.AgentRequest{ + ParentTaskID: "", OwnerSessionID: owner, NodeKey: "", Prompt: "", + Options: workflow.AgentOptions{ + NodeKey: "", AgentName: "", Model: "", Provider: "", ConcurrencyKey: "", Depth: 0, + }, + InvocationIndex: 0, + } +} + +func assertOopsCode(t *testing.T, err error, want string) { + t.Helper() + require.Error(t, err) + + var oopsErr oops.OopsError + require.ErrorAs(t, err, &oopsErr) + assert.Equal(t, want, oopsErr.Code()) +} diff --git a/internal/assistant/workflow_tool.go b/internal/assistant/workflow_tool.go new file mode 100644 index 00000000..ebf292a7 --- /dev/null +++ b/internal/assistant/workflow_tool.go @@ -0,0 +1,130 @@ +package assistant + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "unicode/utf8" + + "github.com/samber/oops" + + "github.com/omarluq/librecode/internal/database" + "github.com/omarluq/librecode/internal/tool" + "github.com/omarluq/librecode/internal/workflow" +) + +const workflowToolName tool.Name = "workflow" + +const workflowTaskIDKey = "workflow_task_id" + +type workflowSubmitter interface { + Submit(context.Context, *workflow.ServiceRequest) (*database.WorkflowRunEntity, error) +} + +type workflowToolExecutor struct { + submitter workflowSubmitter + ownerSessionID string +} + +type workflowToolInput struct { + Arguments map[string]any `json:"arguments,omitempty"` + Name string `json:"name,omitempty"` + Source string `json:"source"` +} + +func (executor *workflowToolExecutor) Definition() tool.Definition { + const rawSchema = `{ + "type":"object", + "additionalProperties":false, + "properties":{ + "name":{"type":"string","description":"Concise workflow name."}, + "source":{"type":"string","description":"MVM Go source importing librecode/workflow."}, + "arguments":{"type":"object","description":"Values exposed as workflow.Arguments."} + }, + "required":["name","source"] + }` + + return tool.Definition{ + Schema: mustToolSchema(rawSchema), Name: workflowToolName, Label: "Start workflow", + Description: "Launch a dynamic, durable asynchronous workflow authored as MVM Go code and return its run ID. " + + "Use this tool when the user asks to launch, start, run, or create a dynamic workflow " + + "or requests durable multi-agent orchestration.", + PromptSnippet: "Launch dynamic workflows and durable multi-agent orchestration", + PromptGuidelines: []string{ + "Treat requests to launch, start, run, or create a dynamic workflow as explicit instructions " + + "to call this tool.", + `Write MVM Go source that imports "librecode/workflow".`, + "Use workflow.Agent, workflow.Wait, workflow.List, workflow.Cancel, or workflow.Pipeline.", + "Use workflows for staged or parallel subagent work; the run continues independently of the model turn.", + }, + ReadOnly: false, + } +} + +func (executor *workflowToolExecutor) Execute(ctx context.Context, input tool.Arguments) (tool.Result, error) { + if executor.submitter == nil { + return tool.TextResult("", nil), oops.In("assistant").Code("workflow_service_unavailable"). + Errorf("workflow service is unavailable") + } + + var args workflowToolInput + if err := input.Decode(&args); err != nil { + return tool.TextResult("", nil), oops.In("assistant").Code("workflow_input").Wrapf(err, "decode workflow input") + } + + args.Name = strings.TrimSpace(args.Name) + args.Source = strings.TrimSpace(args.Source) + + if args.Name == "" { + return tool.TextResult("", nil), oops.In("assistant").Code("workflow_name_required"). + Errorf("workflow name is required") + } + + if args.Source == "" { + return tool.TextResult("", nil), oops.In("assistant").Code("workflow_source_required"). + Errorf("workflow source is required") + } + + if !utf8.ValidString(args.Source) { + return tool.TextResult("", nil), oops.In("assistant").Code("workflow_source_invalid_utf8"). + Errorf("workflow source must be valid UTF-8") + } + + arguments := args.Arguments + if arguments == nil { + arguments = map[string]any{} + } + + argumentsJSON, err := json.Marshal(arguments) + if err != nil { + return tool.TextResult("", nil), oops.In("assistant").Code("encode_workflow_arguments"). + Wrapf(err, "encode workflow arguments") + } + + run, err := executor.submitter.Submit(ctx, &workflow.ServiceRequest{ + Name: args.Name, Source: args.Source, SourceVersion: "v1", + ArgumentsJSON: string(argumentsJSON), OwnerSessionID: executor.ownerSessionID, + }) + if err != nil { + return tool.TextResult("", nil), oops.In("assistant").Code("submit_workflow").Wrapf(err, "submit workflow") + } + + return tool.TextResult( + fmt.Sprintf("Started workflow %q with run ID %s.", run.Name, run.Task.ID), + workflowResultDetails(run), + ), nil +} + +func workflowResultDetails(run *database.WorkflowRunEntity) map[string]any { + if run == nil { + return map[string]any{} + } + + return map[string]any{ + "run_id": run.Task.ID, workflowTaskIDKey: run.Task.ID, + "kind": database.TaskKindWorkflow, executeNameKey: run.Name, "state": run.Task.State, + } +} + +var _ tool.Executor = (*workflowToolExecutor)(nil) diff --git a/internal/assistant/workflow_tool_internal_test.go b/internal/assistant/workflow_tool_internal_test.go new file mode 100644 index 00000000..69c5ed32 --- /dev/null +++ b/internal/assistant/workflow_tool_internal_test.go @@ -0,0 +1,176 @@ +package assistant + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/omarluq/librecode/internal/database" + "github.com/omarluq/librecode/internal/tool" + "github.com/omarluq/librecode/internal/workflow" +) + +const workflowTestSessionID = "workflow-session" + +type workflowSubmitterStub struct { + request *workflow.ServiceRequest + run *database.WorkflowRunEntity + err error +} + +func (stub *workflowSubmitterStub) Submit( + _ context.Context, + request *workflow.ServiceRequest, +) (*database.WorkflowRunEntity, error) { + stub.request = request + + return stub.run, stub.err +} + +func TestWorkflowToolDefinitionSignalsDynamicWorkflowIntent(t *testing.T) { + t.Parallel() + + executor := &workflowToolExecutor{submitter: nil, ownerSessionID: ""} + definition := executor.Definition() + + assert.Contains(t, definition.Description, "launch, start, run, or create a dynamic workflow") + assert.Contains(t, definition.PromptSnippet, "dynamic workflows") + require.NotEmpty(t, definition.PromptGuidelines) + assert.Contains(t, definition.PromptGuidelines[0], "explicit instructions to call this tool") +} + +func TestWorkflowToolSubmitsModelAuthoredSource(t *testing.T) { + t.Parallel() + + stub := &workflowSubmitterStub{ + request: nil, + run: &database.WorkflowRunEntity{ + Task: database.TaskEntity{ + CreatedAt: time.Time{}, StartedAt: nil, FinishedAt: nil, UpdatedAt: time.Time{}, + LeaseExpiresAt: nil, ID: "run-1", Kind: database.TaskKindWorkflow, ParentTaskID: "", + OwnerSessionID: workflowTestSessionID, ConcurrencyKey: "", LeaseOwner: "", + State: database.TaskSucceeded, Result: "reviewed", ErrorCode: "", ErrorMessage: "", + }, + Name: "review", Source: "", SourceHash: "", SourceVersion: "", ArgumentsJSON: "", + }, + err: nil, + } + executor := &workflowToolExecutor{submitter: stub, ownerSessionID: workflowTestSessionID} + input, err := tool.ArgumentsFromRaw([]byte(`{ + "name":"review", + "source":"import \"librecode/workflow\"; workflow.List()", + "arguments":{"scope":"changes"} + }`)) + require.NoError(t, err) + + result, err := executor.Execute(t.Context(), input) + require.NoError(t, err) + assert.Equal(t, `Started workflow "review" with run ID run-1.`, result.Text()) + assert.Equal(t, "run-1", result.Details["run_id"]) + assert.Equal(t, "review", result.Details["name"]) + require.NotNil(t, stub.request) + assert.Equal(t, workflowTestSessionID, stub.request.OwnerSessionID) + assert.Equal(t, "review", stub.request.Name) + assert.JSONEq(t, `{"scope":"changes"}`, stub.request.ArgumentsJSON) +} + +func TestWorkflowToolReturnsWithoutAwaitingCompletion(t *testing.T) { + t.Parallel() + + stub := &workflowSubmitterStub{ + request: nil, + run: &database.WorkflowRunEntity{ + Task: database.TaskEntity{ + CreatedAt: time.Time{}, StartedAt: nil, FinishedAt: nil, UpdatedAt: time.Time{}, + LeaseExpiresAt: nil, ID: "run-queued", Kind: database.TaskKindWorkflow, ParentTaskID: "", + OwnerSessionID: workflowTestSessionID, ConcurrencyKey: "", LeaseOwner: "", + State: database.TaskQueued, Result: "", ErrorCode: "", ErrorMessage: "", + }, + Name: "background review", Source: "", SourceHash: "", SourceVersion: "", ArgumentsJSON: "", + }, + err: nil, + } + executor := &workflowToolExecutor{submitter: stub, ownerSessionID: workflowTestSessionID} + input, err := tool.ArgumentsFromRaw([]byte(`{"name":"background review","source":"1 + 1"}`)) + require.NoError(t, err) + + result, err := executor.Execute(t.Context(), input) + require.NoError(t, err) + assert.Equal(t, `Started workflow "background review" with run ID run-queued.`, result.Text()) + assert.Equal(t, database.TaskQueued, result.Details["state"]) +} + +func TestWorkflowToolRejectsInvalidInput(t *testing.T) { + t.Parallel() + + validStub := &workflowSubmitterStub{request: nil, run: nil, err: nil} + + tests := []struct { + executor *workflowToolExecutor + raw string + name string + want string + }{ + { + name: "missing submitter", executor: &workflowToolExecutor{submitter: nil, ownerSessionID: ""}, + raw: `{}`, want: "workflow service is unavailable", + }, + { + name: "invalid shape", executor: &workflowToolExecutor{submitter: validStub, ownerSessionID: ""}, + raw: `{"name":1}`, want: "decode workflow input", + }, + { + name: "blank name", executor: &workflowToolExecutor{submitter: validStub, ownerSessionID: ""}, + raw: `{"name":" ","source":"1"}`, want: "workflow name is required", + }, + { + name: "blank source", executor: &workflowToolExecutor{submitter: validStub, ownerSessionID: ""}, + raw: `{"name":"run","source":" "}`, want: "workflow source is required", + }, + { + name: "submit failure", + executor: &workflowToolExecutor{ + submitter: &workflowSubmitterStub{request: nil, run: nil, err: errors.New("down")}, + ownerSessionID: "", + }, + raw: `{"name":"run","source":"1"}`, want: "submit workflow", + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + input, err := tool.ArgumentsFromRaw([]byte(test.raw)) + require.NoError(t, err) + _, err = test.executor.Execute(t.Context(), input) + require.ErrorContains(t, err, test.want) + }) + } +} + +func TestWorkflowResultDetailsHandlesNil(t *testing.T) { + t.Parallel() + + assert.Empty(t, workflowResultDetails(nil)) +} + +func TestPromptRegistryIncludesWorkflowWhenConfigured(t *testing.T) { + t.Parallel() + + runtime := new(Runtime) + runtime.profile = topLevelExecutionProfile() + runtime.SetWorkflowSubmitter(&workflowSubmitterStub{request: nil, run: nil, err: nil}) + + registry, err := runtime.promptToolRegistry(t.Context(), t.TempDir(), "owner") + require.NoError(t, err) + assert.True(t, registry.Has(workflowToolName)) + + runtime.SetWorkflowSubmitter(nil) + registry, err = runtime.promptToolRegistry(t.Context(), t.TempDir(), "owner") + require.NoError(t, err) + assert.False(t, registry.Has(workflowToolName)) +} diff --git a/internal/database/agent_task_repository.go b/internal/database/agent_task_repository.go index 4424e1aa..cdabd2ce 100644 --- a/internal/database/agent_task_repository.go +++ b/internal/database/agent_task_repository.go @@ -5,6 +5,7 @@ import ( "database/sql" "encoding/json" "errors" + "time" "github.com/samber/oops" "github.com/vingarcia/ksql" @@ -45,19 +46,98 @@ func NewAgentTaskRepositoryWithProvider(provider ksql.Provider) *AgentTaskReposi return &AgentTaskRepository{sql: provider, tasks: NewTaskRepositoryWithProvider(provider)} } +// CreateWithChildSession atomically creates a child session and its queued agent task. +func (repository *AgentTaskRepository) CreateWithChildSession( + ctx context.Context, + agentTask *AgentTaskEntity, + childRequest *ChildSessionRequest, +) (*AgentTaskEntity, error) { + candidate, child, err := prepareAgentTaskChild(repository.tasks.now(), agentTask, childRequest) + if err != nil { + return nil, err + } + + created, now, err := repository.prepareCreate(candidate) + if err != nil { + return nil, err + } + + if err := repository.sql.Transaction(ctx, func(transaction ksql.Provider) error { + if err := insertSession(ctx, transaction, child); err != nil { + return err + } + + return insertAgentTask(ctx, transaction, created, now) + }); err != nil { + return nil, oops.In("database").Code("create_agent_task_with_session"). + Wrapf(err, "create agent task with child session") + } + + return created, nil +} + // Create persists a generic queued task, its initial event, and agent extension atomically. func (repository *AgentTaskRepository) Create( ctx context.Context, agentTask *AgentTaskEntity, ) (*AgentTaskEntity, error) { + created, now, err := repository.prepareCreate(agentTask) + if err != nil { + return nil, err + } + + if err := repository.sql.Transaction(ctx, func(transaction ksql.Provider) error { + return insertAgentTask(ctx, transaction, created, now) + }); err != nil { + return nil, oops.In("database").Code("create_agent_task").Wrapf(err, "create agent task") + } + + return created, nil +} + +func prepareAgentTaskChild( + now time.Time, + agentTask *AgentTaskEntity, + childRequest *ChildSessionRequest, +) (*AgentTaskEntity, *SessionEntity, error) { + if agentTask == nil { + return nil, nil, oops.In("database").Code("nil_agent_task"). + Errorf("agent task is required") + } + + if childRequest == nil { + return nil, nil, errors.New("database: child session request is required") + } + + if agentTask.Task.OwnerSessionID != childRequest.ParentSessionID { + return nil, nil, errors.New("database: child session parent differs from agent task owner") + } + + child, err := prepareSession(now, childRequest.CWD, childRequest.Name, childRequest.ParentSessionID) + if err != nil { + return nil, nil, err + } + + candidate := *agentTask + candidate.ChildSessionID = child.ID + + return &candidate, child, nil +} + +func (repository *AgentTaskRepository) prepareCreate(agentTask *AgentTaskEntity) (*AgentTaskEntity, time.Time, error) { + if agentTask == nil { + return nil, time.Time{}, oops.In("database").Code("nil_agent_task"). + Errorf("agent task is required") + } + now := repository.tasks.now().UTC() created := *agentTask created.Task.ID = newUUIDv7() created.Task.Kind = TaskKindAgent created.Task.State = TaskQueued created.Task.CreatedAt = now - created.Task.UpdatedAt = now + if created.PolicyJSON == "" { created.PolicyJSON = "{}" } @@ -67,37 +147,30 @@ func (repository *AgentTaskRepository) Create( } if err := validateAgentTaskEntity(&created); err != nil { - return nil, oops.In("database").Code("validate_agent_task").Wrapf(err, "validate agent task") + return nil, time.Time{}, oops.In("database").Code("validate_agent_task").Wrapf(err, "validate agent task") } - if err := repository.sql.Transaction(ctx, func(transaction ksql.Provider) error { - if err := insertTask(ctx, transaction, &created.Task); err != nil { - return err - } + return &created, now, nil +} - const statement = ` -INSERT INTO agent_tasks ( +func insertAgentTask(ctx context.Context, transaction ksql.Provider, created *AgentTaskEntity, now time.Time) error { + if err := insertTask(ctx, transaction, &created.Task); err != nil { + return err + } + + const statement = `INSERT INTO agent_tasks ( task_id, child_session_id, agent_name, prompt, model, provider, policy_json, usage_json, depth ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)` - - _, err := transaction.Exec( - ctx, statement, created.Task.ID, created.ChildSessionID, - created.AgentName, created.Prompt, created.Model, created.Provider, - created.PolicyJSON, created.UsageJSON, created.Depth, - ) - if err != nil { - return oops.In("database").Code("insert_agent_task").Wrapf(err, "insert agent task") - } - - _, err = insertTaskEvent(ctx, transaction, created.Task.ID, 1, "task_queued", "{}", now) - - return err - }); err != nil { - return nil, oops.In("database").Code("create_agent_task").Wrapf(err, "create agent task") + if _, err := transaction.Exec(ctx, statement, created.Task.ID, created.ChildSessionID, + created.AgentName, created.Prompt, created.Model, created.Provider, + created.PolicyJSON, created.UsageJSON, created.Depth); err != nil { + return oops.In("database").Code("insert_agent_task").Wrapf(err, "insert agent task") } - return &created, nil + _, err := insertTaskEvent(ctx, transaction, created.Task.ID, 1, "task_queued", "{}", now) + + return err } // Finish atomically records agent usage, terminal task state, and its event. diff --git a/internal/database/agent_task_repository_test.go b/internal/database/agent_task_repository_test.go index b275b159..2dc407aa 100644 --- a/internal/database/agent_task_repository_test.go +++ b/internal/database/agent_task_repository_test.go @@ -69,6 +69,56 @@ func TestAgentTaskRepositoryLifecycle(t *testing.T) { assert.NotEmpty(t, events[2].Event.ID) } +func TestAgentTaskRepositoryCreatesChildSessionAtomically(t *testing.T) { + t.Parallel() + + fixture := newTaskTestFixture(t) + ctx := t.Context() + owner := fixture.createOwner(ctx) + + candidate := newAgentTask(owner.ID, "") + created, err := fixture.agents.CreateWithChildSession(ctx, candidate, &database.ChildSessionRequest{ + CWD: owner.CWD, Name: "child", ParentSessionID: owner.ID, + }) + require.NoError(t, err) + assert.NotEmpty(t, created.ChildSessionID) + + child, found, err := fixture.sessions.GetSession(ctx, created.ChildSessionID) + require.NoError(t, err) + require.True(t, found) + assert.Equal(t, owner.ID, child.ParentSession) + + mismatched := newAgentTask(owner.ID, "") + _, err = fixture.agents.CreateWithChildSession(ctx, mismatched, &database.ChildSessionRequest{ + CWD: owner.CWD, Name: "wrong-parent", ParentSessionID: testUUIDV7(t), + }) + require.ErrorContains(t, err, "parent differs from agent task owner") + + children, err := fixture.sessions.ListChildSessions(ctx, owner.ID) + require.NoError(t, err) + assert.Len(t, children, 1) +} + +func TestAgentTaskRepositoryRollsBackChildSessionWhenTaskInsertFails(t *testing.T) { + t.Parallel() + + fixture := newTaskTestFixture(t) + ctx := t.Context() + owner := fixture.createOwner(ctx) + candidate := newAgentTask(owner.ID, "") + candidate.Task.ParentTaskID = testUUIDV7(t) + + created, err := fixture.agents.CreateWithChildSession(ctx, candidate, &database.ChildSessionRequest{ + CWD: owner.CWD, Name: "rolled-back-child", ParentSessionID: owner.ID, + }) + require.ErrorContains(t, err, "insert task") + assert.Nil(t, created) + + children, err := fixture.sessions.ListChildSessions(ctx, owner.ID) + require.NoError(t, err) + assert.Empty(t, children, "task insertion failure must roll back the previously inserted child session") +} + func TestAgentTaskRepositoryFinishBehavior(t *testing.T) { t.Parallel() @@ -123,6 +173,11 @@ func TestAgentTaskRepositoryFinishBehavior(t *testing.T) { func TestAgentTaskRepositoryCreateValidationAndDefaults(t *testing.T) { t.Parallel() + fixture := newTaskTestFixture(t) + created, err := fixture.agents.Create(t.Context(), nil) + require.ErrorContains(t, err, "agent task is required") + assert.Nil(t, created) + tests := []struct { name string mutate func(*database.AgentTaskEntity) diff --git a/internal/database/migrations/00008_create_workflow_runs.sql b/internal/database/migrations/00008_create_workflow_runs.sql new file mode 100644 index 00000000..7234389e --- /dev/null +++ b/internal/database/migrations/00008_create_workflow_runs.sql @@ -0,0 +1,28 @@ +-- +goose Up +CREATE TABLE workflow_runs ( + task_id TEXT PRIMARY KEY, + source TEXT NOT NULL, + source_hash TEXT NOT NULL, + source_version TEXT NOT NULL DEFAULT '', + arguments_json TEXT NOT NULL DEFAULT '{}', + FOREIGN KEY (task_id) REFERENCES tasks(id) ON DELETE CASCADE +); + +CREATE TABLE workflow_agent_tasks ( + workflow_task_id TEXT NOT NULL, + agent_task_id TEXT NOT NULL UNIQUE, + sequence INTEGER NOT NULL, + node_key TEXT NOT NULL DEFAULT '', + created_at TEXT NOT NULL, + PRIMARY KEY (workflow_task_id, agent_task_id), + UNIQUE (workflow_task_id, sequence), + FOREIGN KEY (workflow_task_id) REFERENCES workflow_runs(task_id) ON DELETE CASCADE, + FOREIGN KEY (agent_task_id) REFERENCES agent_tasks(task_id) ON DELETE CASCADE +); + +CREATE INDEX idx_workflow_agent_tasks_replay + ON workflow_agent_tasks(workflow_task_id, sequence); + +-- +goose Down +DROP TABLE IF EXISTS workflow_agent_tasks; +DROP TABLE IF EXISTS workflow_runs; diff --git a/internal/database/migrations/00009_add_workflow_invocation_index.sql b/internal/database/migrations/00009_add_workflow_invocation_index.sql new file mode 100644 index 00000000..5b887e03 --- /dev/null +++ b/internal/database/migrations/00009_add_workflow_invocation_index.sql @@ -0,0 +1,63 @@ +-- +goose Up +CREATE TABLE workflow_agent_tasks_v9 ( + workflow_task_id TEXT NOT NULL, + agent_task_id TEXT NOT NULL UNIQUE, + sequence INTEGER NOT NULL, + node_key TEXT NOT NULL DEFAULT '', + invocation_index INTEGER NOT NULL DEFAULT 0, + created_at TEXT NOT NULL, + PRIMARY KEY (workflow_task_id, agent_task_id), + UNIQUE (workflow_task_id, sequence), + UNIQUE (workflow_task_id, node_key, invocation_index), + FOREIGN KEY (workflow_task_id) REFERENCES workflow_runs(task_id) ON DELETE CASCADE, + FOREIGN KEY (agent_task_id) REFERENCES agent_tasks(task_id) ON DELETE CASCADE, + CHECK (sequence > 0), + CHECK (invocation_index >= 0) +); + +INSERT INTO workflow_agent_tasks_v9 + (workflow_task_id, agent_task_id, sequence, node_key, invocation_index, created_at) +SELECT current.workflow_task_id, + current.agent_task_id, + current.sequence, + current.node_key, + ( + SELECT COUNT(*) - 1 + FROM workflow_agent_tasks AS preceding + WHERE preceding.workflow_task_id = current.workflow_task_id + AND preceding.node_key = current.node_key + AND preceding.sequence <= current.sequence + ), + current.created_at +FROM workflow_agent_tasks AS current; + +DROP TABLE workflow_agent_tasks; +ALTER TABLE workflow_agent_tasks_v9 RENAME TO workflow_agent_tasks; + +CREATE INDEX idx_workflow_agent_tasks_replay + ON workflow_agent_tasks(workflow_task_id, sequence); + +-- +goose Down +CREATE TABLE workflow_agent_tasks_v8 ( + workflow_task_id TEXT NOT NULL, + agent_task_id TEXT NOT NULL UNIQUE, + sequence INTEGER NOT NULL, + node_key TEXT NOT NULL DEFAULT '', + created_at TEXT NOT NULL, + PRIMARY KEY (workflow_task_id, agent_task_id), + UNIQUE (workflow_task_id, sequence), + FOREIGN KEY (workflow_task_id) REFERENCES workflow_runs(task_id) ON DELETE CASCADE, + FOREIGN KEY (agent_task_id) REFERENCES agent_tasks(task_id) ON DELETE CASCADE, + CHECK (sequence > 0) +); + +INSERT INTO workflow_agent_tasks_v8 + (workflow_task_id, agent_task_id, sequence, node_key, created_at) +SELECT workflow_task_id, agent_task_id, sequence, node_key, created_at +FROM workflow_agent_tasks; + +DROP TABLE workflow_agent_tasks; +ALTER TABLE workflow_agent_tasks_v8 RENAME TO workflow_agent_tasks; + +CREATE INDEX idx_workflow_agent_tasks_replay + ON workflow_agent_tasks(workflow_task_id, sequence); diff --git a/internal/database/migrations/00010_add_workflow_name.sql b/internal/database/migrations/00010_add_workflow_name.sql new file mode 100644 index 00000000..b7797fa8 --- /dev/null +++ b/internal/database/migrations/00010_add_workflow_name.sql @@ -0,0 +1,7 @@ +-- +goose Up +ALTER TABLE workflow_runs + ADD COLUMN name TEXT NOT NULL DEFAULT ''; + +-- +goose Down +ALTER TABLE workflow_runs + DROP COLUMN name; diff --git a/internal/database/migrations_test.go b/internal/database/migrations_test.go new file mode 100644 index 00000000..6d2bd0ad --- /dev/null +++ b/internal/database/migrations_test.go @@ -0,0 +1,167 @@ +package database_test + +import ( + "context" + "database/sql" + "testing" + "testing/fstest" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/omarluq/librecode/internal/database" +) + +const deployedWorkflowMigrationV8 = `-- +goose Up +CREATE TABLE workflow_runs ( + task_id TEXT PRIMARY KEY, + source TEXT NOT NULL, + source_hash TEXT NOT NULL, + source_version TEXT NOT NULL DEFAULT '', + arguments_json TEXT NOT NULL DEFAULT '{}', + FOREIGN KEY (task_id) REFERENCES tasks(id) ON DELETE CASCADE +); +CREATE TABLE workflow_agent_tasks ( + workflow_task_id TEXT NOT NULL, + agent_task_id TEXT NOT NULL UNIQUE, + sequence INTEGER NOT NULL, + node_key TEXT NOT NULL DEFAULT '', + created_at TEXT NOT NULL, + PRIMARY KEY (workflow_task_id, agent_task_id), + UNIQUE (workflow_task_id, sequence), + FOREIGN KEY (workflow_task_id) REFERENCES workflow_runs(task_id) ON DELETE CASCADE, + FOREIGN KEY (agent_task_id) REFERENCES agent_tasks(task_id) ON DELETE CASCADE +); +CREATE INDEX idx_workflow_agent_tasks_replay + ON workflow_agent_tasks(workflow_task_id, sequence); +-- +goose Down +DROP TABLE IF EXISTS workflow_agent_tasks; +DROP TABLE IF EXISTS workflow_runs; +` + +func TestMigrateAddsWorkflowInvocationIndexToDeployedVersionEightDatabase(t *testing.T) { + t.Parallel() + + connection := newMigratedThroughVersion(t, 7) + ctx := context.Background() + migrationRoot := fstest.MapFS{ + "00008_create_workflow_runs.sql": &fstest.MapFile{Data: []byte(deployedWorkflowMigrationV8)}, + } + provider, err := database.NewMigrationProvider(connection, migrationRoot) + require.NoError(t, err) + + _, err = provider.Up(ctx) + require.NoError(t, err) + + insertWorkflowMigrationFixtures(ctx, t, connection) + require.NoError(t, database.Migrate(ctx, connection)) + + assertWorkflowInvocationColumn(ctx, t, connection) + assertWorkflowInvocationIndexes(ctx, t, connection) +} + +func insertWorkflowMigrationFixtures(ctx context.Context, t *testing.T, connection *sql.DB) { + t.Helper() + + _, err := connection.ExecContext(ctx, ` +INSERT INTO sessions (id, cwd, name, created_at, updated_at) +VALUES ('01900000-0000-7000-8000-000000000001', '/tmp', 'owner', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP), + ('01900000-0000-7000-8000-000000000004', '/tmp', 'child', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP), + ('01900000-0000-7000-8000-000000000006', '/tmp', 'second child', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)`) + require.NoError(t, err) + + insertWorkflowTasks(ctx, t, connection) + insertWorkflowAgentLinks(ctx, t, connection) +} + +func insertWorkflowTasks(ctx context.Context, t *testing.T, connection *sql.DB) { + t.Helper() + + const ownerID = "01900000-0000-7000-8000-000000000001" + + _, err := connection.ExecContext(ctx, ` +INSERT INTO tasks (id, kind, state, owner_session_id, created_at, updated_at) +VALUES ('01900000-0000-7000-8000-000000000002', 'workflow', 'queued', ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP), + ('01900000-0000-7000-8000-000000000003', 'agent', 'queued', ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP), + ('01900000-0000-7000-8000-000000000005', 'agent', 'queued', ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)`, + ownerID, ownerID, ownerID) + require.NoError(t, err) +} + +func insertWorkflowAgentLinks(ctx context.Context, t *testing.T, connection *sql.DB) { + t.Helper() + + _, err := connection.ExecContext(ctx, ` +INSERT INTO workflow_runs (task_id, source, source_hash) +VALUES ('01900000-0000-7000-8000-000000000002', 'package main', 'hash')`) + require.NoError(t, err) + + _, err = connection.ExecContext(ctx, ` +INSERT INTO agent_tasks (task_id, child_session_id, agent_name, prompt, depth) +VALUES ('01900000-0000-7000-8000-000000000003', '01900000-0000-7000-8000-000000000004', 'reviewer', 'review', 1), + ('01900000-0000-7000-8000-000000000005', '01900000-0000-7000-8000-000000000006', 'reviewer', 'review again', 1)`) + require.NoError(t, err) + + _, err = connection.ExecContext(ctx, ` +INSERT INTO workflow_agent_tasks + (workflow_task_id, agent_task_id, sequence, node_key, created_at) +VALUES + ('01900000-0000-7000-8000-000000000002', '01900000-0000-7000-8000-000000000003', 1, 'review', CURRENT_TIMESTAMP), + ('01900000-0000-7000-8000-000000000002', '01900000-0000-7000-8000-000000000005', 2, 'review', CURRENT_TIMESTAMP)`) + require.NoError(t, err) +} + +func assertWorkflowInvocationColumn(ctx context.Context, t *testing.T, connection *sql.DB) { + t.Helper() + + columns, err := connection.QueryContext(ctx, `PRAGMA table_info(workflow_agent_tasks)`) + require.NoError(t, err) + + t.Cleanup(func() { require.NoError(t, columns.Close()) }) + + found := false + + for columns.Next() { + var ( + cid, notNull, primaryKey int + name, columnType string + defaultValue sql.NullString + ) + + require.NoError(t, columns.Scan(&cid, &name, &columnType, ¬Null, &defaultValue, &primaryKey)) + + if name != "invocation_index" { + continue + } + + found = true + + assert.Equal(t, 1, notNull) + assert.Equal(t, "0", defaultValue.String) + } + + require.NoError(t, columns.Err()) + require.True(t, found) +} + +func assertWorkflowInvocationIndexes(ctx context.Context, t *testing.T, connection *sql.DB) { + t.Helper() + + rows, err := connection.QueryContext(ctx, + `SELECT invocation_index FROM workflow_agent_tasks ORDER BY sequence`) + require.NoError(t, err) + + t.Cleanup(func() { require.NoError(t, rows.Close()) }) + + var invocationIndexes []int + + for rows.Next() { + var invocationIndex int + + require.NoError(t, rows.Scan(&invocationIndex)) + invocationIndexes = append(invocationIndexes, invocationIndex) + } + + require.NoError(t, rows.Err()) + assert.Equal(t, []int{0, 1}, invocationIndexes) +} diff --git a/internal/database/repository_helpers.go b/internal/database/repository_helpers.go index 129be096..4bddc6c4 100644 --- a/internal/database/repository_helpers.go +++ b/internal/database/repository_helpers.go @@ -1,12 +1,62 @@ package database import ( + "context" + "errors" "fmt" "time" "github.com/samber/lo" + "github.com/samber/oops" + "github.com/vingarcia/ksql" ) +func querySQLRow[Row any, Entity any]( + ctx context.Context, + provider ksql.Provider, + convert func(*Row) (*Entity, error), + query string, + entity string, + args ...any, +) (entityValue *Entity, found bool, queryErr error) { + var row Row + if err := provider.QueryOne(ctx, &row, query, args...); err != nil { + if errors.Is(err, ksql.ErrRecordNotFound) { + return nil, false, nil + } + + return nil, false, oops.In("database").Code("query_"+entity).Wrapf(err, "query %s", entity) + } + + converted, err := convert(&row) + if err != nil { + return nil, false, oops.In("database").Code("scan_"+entity).Wrapf(err, "scan %s", entity) + } + + return converted, true, nil +} + +func querySQLRows[Row any, Entity any]( + ctx context.Context, + provider ksql.Provider, + convert func(*Row) (*Entity, error), + query string, + entity string, + args ...any, +) ([]Entity, error) { + rows := []Row{} + if err := provider.Query(ctx, &rows, query, args...); err != nil { + return nil, oops.In("database").Code("query_"+entity).Wrapf(err, "query %s", entity) + } + + entities, err := collectSQLRows(rows, convert) + if err != nil { + return nil, oops.In("database").Code("scan_"+entity).Wrapf(err, "scan %s", entity) + } + + return entities, nil +} + func collectSQLRows[T any, R any](rows []T, convert func(*T) (*R, error)) ([]R, error) { output, err := lo.MapErr(rows, func(row T, _ int) (R, error) { item, err := convert(&row) diff --git a/internal/database/session_repository.go b/internal/database/session_repository.go index ce4ec013..bfea72df 100644 --- a/internal/database/session_repository.go +++ b/internal/database/session_repository.go @@ -12,6 +12,13 @@ import ( ksqlite "github.com/vingarcia/ksql/adapters/modernc-ksqlite" ) +// ChildSessionRequest describes a child session created with durable agent work. +type ChildSessionRequest struct { + CWD string + Name string + ParentSessionID string +} + // SessionRepository provides persistence for sessions and tree entries. type SessionRepository struct { sql ksql.Provider @@ -74,45 +81,60 @@ func newSessionID() string { return newUUIDv7() } -// CreateSession creates a new persisted session for a working directory. -func (repository *SessionRepository) CreateSession( - ctx context.Context, - cwd string, - name string, - parentSession string, -) (*SessionEntity, error) { - timestamp := repository.now().UTC() - createdSession := SessionEntity{ - CreatedAt: timestamp, - UpdatedAt: timestamp, +func prepareSession(now time.Time, cwd, name, parentSession string) (*SessionEntity, error) { + created := &SessionEntity{ + CreatedAt: now.UTC(), + UpdatedAt: now.UTC(), ID: newSessionID(), CWD: cwd, Name: name, ParentSession: parentSession, } + if err := validateSessionEntity(created); err != nil { + return nil, oops.In("database").Code("validate_session").Wrapf(err, "validate session") + } + + return created, nil +} +func insertSession(ctx context.Context, provider ksql.Provider, session *SessionEntity) error { const statement = ` INSERT INTO sessions (id, cwd, name, parent_session, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)` - if err := validateSessionEntity(&createdSession); err != nil { - return nil, oops.In("database").Code("validate_session").Wrapf(err, "validate session") - } - - if _, err := repository.sql.Exec( + if _, err := provider.Exec( ctx, statement, - createdSession.ID, - createdSession.CWD, - createdSession.Name, - createdSession.ParentSession, - formatTime(createdSession.CreatedAt), - formatTime(createdSession.UpdatedAt), + session.ID, + session.CWD, + session.Name, + session.ParentSession, + formatTime(session.CreatedAt), + formatTime(session.UpdatedAt), ); err != nil { - return nil, oops.In("database").Code("create_session").Wrapf(err, "create session") + return oops.In("database").Code("create_session").Wrapf(err, "create session") + } + + return nil +} + +// CreateSession creates a new persisted session for a working directory. +func (repository *SessionRepository) CreateSession( + ctx context.Context, + cwd string, + name string, + parentSession string, +) (*SessionEntity, error) { + created, err := prepareSession(repository.now(), cwd, name, parentSession) + if err != nil { + return nil, err + } + + if err := insertSession(ctx, repository.sql, created); err != nil { + return nil, err } - return &createdSession, nil + return created, nil } // LatestSession returns the newest top-level session for cwd. diff --git a/internal/database/task_repository.go b/internal/database/task_repository.go index ebd6281c..dfdb4e3e 100644 --- a/internal/database/task_repository.go +++ b/internal/database/task_repository.go @@ -157,6 +157,15 @@ func (repository *TaskRepository) Create(ctx context.Context, task *TaskEntity) // ClaimQueued atomically moves a queued task to running, assigns its lease, and appends an event. func (repository *TaskRepository) ClaimQueued(ctx context.Context, claim *TaskClaim) (bool, error) { + return repository.claim(ctx, claim, TaskQueued) +} + +// ClaimInterrupted atomically resumes an interrupted task as running with a new lease. +func (repository *TaskRepository) ClaimInterrupted(ctx context.Context, claim *TaskClaim) (bool, error) { + return repository.claim(ctx, claim, TaskInterrupted) +} + +func (repository *TaskRepository) claim(ctx context.Context, claim *TaskClaim, from TaskState) (bool, error) { if claim == nil || strings.TrimSpace(claim.LeaseOwner) == "" || claim.LeaseExpiresAt.IsZero() { return false, errors.New("database: task claim requires an owner and expiry") } @@ -170,13 +179,14 @@ func (repository *TaskRepository) ClaimQueued(ctx context.Context, claim *TaskCl err := repository.sql.Transaction(ctx, func(transaction ksql.Provider) error { now := repository.now().UTC() - const update = `UPDATE tasks SET state = ?, started_at = COALESCE(started_at, ?), -updated_at = ?, lease_owner = ?, lease_expires_at = ? WHERE id = ? AND state = ?` + const update = `UPDATE tasks SET state = ?, started_at = COALESCE(started_at, ?), finished_at = NULL, +updated_at = ?, lease_owner = ?, lease_expires_at = ?, result = '', error_code = '', error_message = '' +WHERE id = ? AND state = ?` result, err := transaction.Exec(ctx, update, TaskRunning, formatTime(now), formatTime(now), - claim.LeaseOwner, formatTime(claim.LeaseExpiresAt.UTC()), claim.TaskID, TaskQueued) + claim.LeaseOwner, formatTime(claim.LeaseExpiresAt.UTC()), claim.TaskID, from) if err != nil { - return oops.In("database").Code("claim_task").Wrapf(err, "claim queued task") + return oops.In("database").Code("claim_task").Wrapf(err, "claim task") } rows, err := result.RowsAffected() @@ -199,7 +209,7 @@ updated_at = ?, lease_owner = ?, lease_expires_at = ? WHERE id = ? AND state = ? return err }) if err != nil { - return false, oops.In("database").Code("claim_task").Wrapf(err, "claim queued task") + return false, oops.In("database").Code("claim_task").Wrapf(err, "claim task") } return changed, nil diff --git a/internal/database/task_test_helpers_test.go b/internal/database/task_test_helpers_test.go index 049bf607..ade3ea6b 100644 --- a/internal/database/task_test_helpers_test.go +++ b/internal/database/task_test_helpers_test.go @@ -23,10 +23,11 @@ const ( ) type taskTestFixture struct { - t *testing.T - agents *database.AgentTaskRepository - tasks *database.TaskRepository - sessions *database.SessionRepository + t *testing.T + agents *database.AgentTaskRepository + tasks *database.TaskRepository + workflows *database.WorkflowRepository + sessions *database.SessionRepository } func newTaskTestFixture(t *testing.T) *taskTestFixture { @@ -36,10 +37,11 @@ func newTaskTestFixture(t *testing.T) *taskTestFixture { require.NoError(t, database.Migrate(t.Context(), connection)) return &taskTestFixture{ - t: t, - agents: database.NewAgentTaskRepository(connection), - tasks: database.NewTaskRepository(connection), - sessions: database.NewSessionRepository(connection), + t: t, + agents: database.NewAgentTaskRepository(connection), + tasks: database.NewTaskRepository(connection), + workflows: database.NewWorkflowRepository(connection), + sessions: database.NewSessionRepository(connection), } } diff --git a/internal/database/validation.go b/internal/database/validation.go index 36f14116..514e9733 100644 --- a/internal/database/validation.go +++ b/internal/database/validation.go @@ -111,6 +111,31 @@ func validateTaskEntity(entity *TaskEntity) error { return validateRequiredTime("task.updated_at", entity.UpdatedAt) } +func validateWorkflowRunEntity(entity *WorkflowRunEntity) error { + if err := validateTaskEntity(&entity.Task); err != nil { + return err + } + + if entity.Task.Kind != TaskKindWorkflow { + return errors.New("workflow_run.task.kind must be workflow") + } + + if err := validateRequiredText("workflow_run.source", entity.Source); err != nil { + return err + } + + if err := validateRequiredText("workflow_run.source_hash", entity.SourceHash); err != nil { + return err + } + + var arguments map[string]any + if err := json.Unmarshal([]byte(entity.ArgumentsJSON), &arguments); err != nil || arguments == nil { + return errors.New("workflow_run.arguments_json must be a JSON object") + } + + return nil +} + func validateAgentTaskEntity(entity *AgentTaskEntity) error { if err := validateTaskEntity(&entity.Task); err != nil { return err diff --git a/internal/database/workflow_repository.go b/internal/database/workflow_repository.go new file mode 100644 index 00000000..6d8237d4 --- /dev/null +++ b/internal/database/workflow_repository.go @@ -0,0 +1,468 @@ +package database + +import ( + "context" + "database/sql" + "errors" + "strings" + "time" + + "github.com/samber/oops" + "github.com/vingarcia/ksql" + ksqlite "github.com/vingarcia/ksql/adapters/modernc-ksqlite" + "modernc.org/sqlite" + sqlite3 "modernc.org/sqlite/lib" +) + +// TaskKindWorkflow identifies durable workflow execution. +const TaskKindWorkflow = "workflow" + +// WorkflowRunEntity contains workflow-specific data for a generic task. +type WorkflowRunEntity struct { + Task TaskEntity + Name string + Source string + SourceHash string + SourceVersion string + ArgumentsJSON string +} + +// WorkflowAgentTaskEntity associates an agent task with its workflow-local launch order. +type WorkflowAgentTaskEntity struct { + CreatedAt time.Time + WorkflowTaskID string + AgentTaskID string + NodeKey string + InvocationIndex int + Sequence int64 +} + +// WorkflowRepository persists workflow metadata and composes generic lifecycle operations. +type WorkflowRepository struct { + sql ksql.Provider + tasks *TaskRepository + agentTasks *AgentTaskRepository +} + +// NewWorkflowRepository creates a workflow repository. +func NewWorkflowRepository(connection *sql.DB) *WorkflowRepository { + provider, err := ksqlite.NewFromSQLDB(connection) + if err != nil { + panic(err) + } + + return NewWorkflowRepositoryWithProvider(provider) +} + +// NewWorkflowRepositoryWithProvider creates a workflow repository with an explicit SQL provider. +func NewWorkflowRepositoryWithProvider(provider ksql.Provider) *WorkflowRepository { + return &WorkflowRepository{ + sql: provider, tasks: NewTaskRepositoryWithProvider(provider), + agentTasks: NewAgentTaskRepositoryWithProvider(provider), + } +} + +// Tasks returns the generic task repository used for workflow lifecycle and events. +func (repository *WorkflowRepository) Tasks() *TaskRepository { + return repository.tasks +} + +// AgentTasks returns the agent-task repository sharing this repository's transaction provider. +func (repository *WorkflowRepository) AgentTasks() *AgentTaskRepository { + return repository.agentTasks +} + +// Create persists a queued workflow task, metadata, and initial event atomically. +func (repository *WorkflowRepository) Create( + ctx context.Context, + run *WorkflowRunEntity, +) (*WorkflowRunEntity, error) { + if run == nil { + return nil, errors.New("database: workflow run is required") + } + + now := repository.tasks.now().UTC() + created := *run + created.Task.ID = newUUIDv7() + created.Task.Kind = TaskKindWorkflow + created.Task.State = TaskQueued + created.Task.CreatedAt = now + + created.Task.UpdatedAt = now + if created.ArgumentsJSON == "" { + created.ArgumentsJSON = "{}" + } + + if err := validateWorkflowRunEntity(&created); err != nil { + return nil, oops.In("database").Code("validate_workflow_run").Wrapf(err, "validate workflow run") + } + + if err := repository.sql.Transaction(ctx, func(transaction ksql.Provider) error { + if err := insertTask(ctx, transaction, &created.Task); err != nil { + return err + } + + const statement = `INSERT INTO workflow_runs +(task_id, name, source, source_hash, source_version, arguments_json) VALUES (?, ?, ?, ?, ?, ?)` + if _, err := transaction.Exec(ctx, statement, created.Task.ID, created.Name, created.Source, created.SourceHash, + created.SourceVersion, created.ArgumentsJSON); err != nil { + return oops.In("database").Code("insert_workflow_run").Wrapf(err, "insert workflow run") + } + + _, err := insertTaskEvent(ctx, transaction, created.Task.ID, 1, "task_queued", "{}", now) + + return err + }); err != nil { + return nil, oops.In("database").Code("create_workflow_run").Wrapf(err, "create workflow run") + } + + return &created, nil +} + +// Get loads a workflow run and its generic lifecycle by task ID. +func (repository *WorkflowRepository) Get( + ctx context.Context, + taskID string, +) (*WorkflowRunEntity, bool, error) { + const query = `SELECT ` + workflowRunColumns + ` FROM tasks t +JOIN workflow_runs w ON w.task_id = t.id WHERE t.id = ? AND t.kind = ?` + + return querySQLRow(ctx, repository.sql, workflowRunFromRow, query, "workflow_run", taskID, TaskKindWorkflow) +} + +// ListByOwner returns workflow runs belonging to a session, newest first. +func (repository *WorkflowRepository) ListByOwner( + ctx context.Context, + ownerSessionID string, + limit int, +) ([]WorkflowRunEntity, error) { + if limit <= 0 { + limit = 100 + } + + const query = `SELECT ` + workflowRunColumns + ` FROM tasks t +JOIN workflow_runs w ON w.task_id = t.id +WHERE t.kind = ? AND t.owner_session_id = ? +ORDER BY t.updated_at DESC, t.id DESC LIMIT ?` + + return querySQLRows( + ctx, repository.sql, workflowRunFromRow, query, "workflow_run", TaskKindWorkflow, ownerSessionID, limit, + ) +} + +// CreateAgentTaskWithChildSession atomically creates a child session, queued agent task, and workflow link. +func (repository *WorkflowRepository) CreateAgentTaskWithChildSession( + ctx context.Context, + workflowTaskID string, + agentTask *AgentTaskEntity, + childRequest *ChildSessionRequest, + nodeKey string, + invocationIndex int, +) (*AgentTaskEntity, error) { + candidate, child, err := prepareAgentTaskChild(repository.tasks.now(), agentTask, childRequest) + if err != nil { + return nil, err + } + + return repository.createAgentTask(ctx, workflowTaskID, candidate, child, nodeKey, invocationIndex) +} + +// CreateAgentTask atomically persists a queued agent task and its workflow link. +func (repository *WorkflowRepository) CreateAgentTask( + ctx context.Context, + workflowTaskID string, + agentTask *AgentTaskEntity, + nodeKey string, + invocationIndex int, +) (*AgentTaskEntity, error) { + return repository.createAgentTask(ctx, workflowTaskID, agentTask, nil, nodeKey, invocationIndex) +} + +func (repository *WorkflowRepository) createAgentTask( + ctx context.Context, + workflowTaskID string, + agentTask *AgentTaskEntity, + child *SessionEntity, + nodeKey string, + invocationIndex int, +) (*AgentTaskEntity, error) { + if err := validateUUIDv7("workflow_agent_task.workflow_task_id", workflowTaskID); err != nil { + return nil, err + } + + agents := NewAgentTaskRepositoryWithProvider(repository.sql) + + created, _, err := agents.prepareCreate(agentTask) + if err != nil { + return nil, err + } + + if created.Task.ParentTaskID != workflowTaskID { + return nil, oops.In("database").Code("workflow_agent_parent_mismatch"). + Errorf("agent task parent must be workflow %q", workflowTaskID) + } + + request := workflowAgentTaskCreate{ + workflowTaskID: workflowTaskID, + agentTask: created, + child: child, + nodeKey: strings.TrimSpace(nodeKey), + invocationIndex: invocationIndex, + } + + if err := repository.sql.Transaction(ctx, func(transaction ksql.Provider) error { + return persistWorkflowAgentTask(ctx, transaction, &request) + }); err != nil { + return nil, oops.In("database").Code("create_workflow_agent_task").Wrapf(err, "create workflow agent task") + } + + return created, nil +} + +type workflowAgentTaskCreate struct { + agentTask *AgentTaskEntity + child *SessionEntity + workflowTaskID string + nodeKey string + + invocationIndex int +} + +func persistWorkflowAgentTask( + ctx context.Context, + transaction ksql.Provider, + request *workflowAgentTaskCreate, +) error { + var workflow struct { + OwnerSessionID string `ksql:"owner_session_id"` + } + + const workflowQuery = `SELECT t.owner_session_id FROM tasks t +JOIN workflow_runs w ON w.task_id = t.id WHERE t.id = ?` + if err := transaction.QueryOne(ctx, &workflow, workflowQuery, request.workflowTaskID); err != nil { + return oops.In("database").Code("load_workflow_agent_parent").Wrapf(err, "load workflow agent parent") + } + + if workflow.OwnerSessionID != request.agentTask.Task.OwnerSessionID { + return oops.In("database").Code("workflow_agent_owner_mismatch"). + Errorf("agent task owner differs from workflow owner") + } + + if request.child != nil { + if err := insertSession(ctx, transaction, request.child); err != nil { + return err + } + } + + if err := insertAgentTask(ctx, transaction, request.agentTask, request.agentTask.Task.CreatedAt); err != nil { + return err + } + + _, err := insertWorkflowAgentTask( + ctx, + transaction, + request.workflowTaskID, + request.agentTask.Task.ID, + request.nodeKey, + request.invocationIndex, + request.agentTask.Task.CreatedAt, + ) + + return err +} + +// LinkAgentTask appends an agent task to a workflow's launch order. Repeating +// the exact link is safe and returns the existing row. +func (repository *WorkflowRepository) LinkAgentTask( + ctx context.Context, + workflowTaskID string, + agentTaskID string, + nodeKey string, + invocationIndex int, +) (*WorkflowAgentTaskEntity, error) { + if err := validateUUIDv7("workflow_agent_task.workflow_task_id", workflowTaskID); err != nil { + return nil, err + } + + if err := validateUUIDv7("workflow_agent_task.agent_task_id", agentTaskID); err != nil { + return nil, err + } + + var link *WorkflowAgentTaskEntity + + nodeKey = strings.TrimSpace(nodeKey) + + err := repository.sql.Transaction(ctx, func(transaction ksql.Provider) error { + created, err := insertWorkflowAgentTask(ctx, transaction, workflowTaskID, agentTaskID, + nodeKey, invocationIndex, repository.tasks.now().UTC()) + if err == nil { + link = created + + return nil + } + + if !isWorkflowInvocationUniqueConstraint(err) { + return err + } + + existing, found, queryErr := querySQLRow(ctx, transaction, workflowAgentTaskFromRow, + `SELECT workflow_task_id, agent_task_id, sequence, node_key, invocation_index, created_at +FROM workflow_agent_tasks WHERE workflow_task_id = ? AND node_key = ? AND invocation_index = ?`, + "workflow_agent_task", workflowTaskID, nodeKey, invocationIndex) + if queryErr != nil || !found { + return err + } + + if existing.AgentTaskID != agentTaskID { + return oops.In("database").Code("workflow_agent_invocation_conflict"). + Errorf("workflow invocation is already linked to agent task %q", existing.AgentTaskID) + } + + link = existing + + return nil + }) + if err != nil { + return nil, oops.In("database").Code("create_workflow_agent_link").Wrapf(err, "create workflow agent link") + } + + return link, nil +} + +func isWorkflowInvocationUniqueConstraint(err error) bool { + var sqliteErr *sqlite.Error + if !errors.As(err, &sqliteErr) || sqliteErr.Code() != sqlite3.SQLITE_CONSTRAINT_UNIQUE { + return false + } + + const constraint = "UNIQUE constraint failed: workflow_agent_tasks.workflow_task_id, " + + "workflow_agent_tasks.node_key, workflow_agent_tasks.invocation_index" + + return strings.Contains(sqliteErr.Error(), constraint) +} + +func insertWorkflowAgentTask( + ctx context.Context, + transaction ksql.Provider, + workflowTaskID string, + agentTaskID string, + nodeKey string, + invocationIndex int, + now time.Time, +) (*WorkflowAgentTaskEntity, error) { + var row struct { + Sequence int64 `ksql:"sequence"` + } + + const statement = `INSERT INTO workflow_agent_tasks +(workflow_task_id, agent_task_id, sequence, node_key, invocation_index, created_at) +SELECT ?, ?, COALESCE(MAX(sequence), 0) + 1, ?, ?, ? +FROM workflow_agent_tasks WHERE workflow_task_id = ? +RETURNING sequence` + if err := transaction.QueryOne(ctx, &row, statement, workflowTaskID, agentTaskID, + nodeKey, invocationIndex, formatTime(now), workflowTaskID); err != nil { + return nil, oops.In("database").Code("link_workflow_agent_task").Wrapf(err, "link workflow agent task") + } + + return &WorkflowAgentTaskEntity{CreatedAt: now, WorkflowTaskID: workflowTaskID, + AgentTaskID: agentTaskID, NodeKey: nodeKey, InvocationIndex: invocationIndex, Sequence: row.Sequence}, nil +} + +// FindAgentTask returns a linked child by its normalized workflow invocation identity. +func (repository *WorkflowRepository) FindAgentTask( + ctx context.Context, + workflowTaskID string, + nodeKey string, + invocationIndex int, +) (*WorkflowAgentTaskEntity, bool, error) { + const query = `SELECT workflow_task_id, agent_task_id, sequence, node_key, invocation_index, created_at +FROM workflow_agent_tasks WHERE workflow_task_id = ? AND node_key = ? AND invocation_index = ?` + + return querySQLRow(ctx, repository.sql, workflowAgentTaskFromRow, query, "workflow_agent_task", + workflowTaskID, strings.TrimSpace(nodeKey), invocationIndex) +} + +// ListAgentTasks returns linked agent tasks in launch order. +func (repository *WorkflowRepository) ListAgentTasks( + ctx context.Context, + workflowTaskID string, +) ([]WorkflowAgentTaskEntity, error) { + const query = `SELECT workflow_task_id, agent_task_id, sequence, node_key, invocation_index, created_at +FROM workflow_agent_tasks WHERE workflow_task_id = ? ORDER BY sequence ASC` + + rows := []workflowAgentTaskRow{} + if err := repository.sql.Query(ctx, &rows, query, workflowTaskID); err != nil { + return nil, oops.In("database").Code("list_workflow_agent_tasks").Wrapf(err, "list workflow agent tasks") + } + + links, err := collectSQLRows(rows, workflowAgentTaskFromRow) + if err != nil { + return nil, oops.In("database").Code("scan_workflow_agent_task").Wrapf(err, "scan workflow agent task") + } + + return links, nil +} + +const workflowRunColumns = `t.id, t.kind, t.parent_task_id, t.owner_session_id, t.concurrency_key, +t.state, t.result, t.error_code, t.error_message, t.created_at, t.started_at, t.finished_at, +t.updated_at, t.lease_owner, t.lease_expires_at, +w.name, w.source, w.source_hash, w.source_version, w.arguments_json` + +type workflowRunRow struct { + ID string `ksql:"id"` + Kind string `ksql:"kind"` + ParentTaskID *string `ksql:"parent_task_id"` + OwnerSessionID string `ksql:"owner_session_id"` + ConcurrencyKey string `ksql:"concurrency_key"` + State string `ksql:"state"` + Result string `ksql:"result"` + ErrorCode string `ksql:"error_code"` + ErrorMessage string `ksql:"error_message"` + CreatedAt string `ksql:"created_at"` + StartedAt *string `ksql:"started_at"` + FinishedAt *string `ksql:"finished_at"` + UpdatedAt string `ksql:"updated_at"` + LeaseOwner *string `ksql:"lease_owner"` + LeaseExpiresAt *string `ksql:"lease_expires_at"` + Name string `ksql:"name"` + Source string `ksql:"source"` + SourceHash string `ksql:"source_hash"` + SourceVersion string `ksql:"source_version"` + ArgumentsJSON string `ksql:"arguments_json"` +} + +func workflowRunFromRow(row *workflowRunRow) (*WorkflowRunEntity, error) { + task, err := taskFromRow(&taskRow{ + ID: row.ID, Kind: row.Kind, ParentTaskID: row.ParentTaskID, OwnerSessionID: row.OwnerSessionID, + ConcurrencyKey: row.ConcurrencyKey, State: row.State, Result: row.Result, ErrorCode: row.ErrorCode, + ErrorMessage: row.ErrorMessage, CreatedAt: row.CreatedAt, StartedAt: row.StartedAt, + FinishedAt: row.FinishedAt, UpdatedAt: row.UpdatedAt, LeaseOwner: row.LeaseOwner, + LeaseExpiresAt: row.LeaseExpiresAt, + }) + if err != nil { + return nil, err + } + + return &WorkflowRunEntity{Task: *task, Name: row.Name, Source: row.Source, SourceHash: row.SourceHash, + SourceVersion: row.SourceVersion, ArgumentsJSON: row.ArgumentsJSON}, nil +} + +type workflowAgentTaskRow struct { + CreatedAt string `ksql:"created_at"` + WorkflowTaskID string `ksql:"workflow_task_id"` + AgentTaskID string `ksql:"agent_task_id"` + NodeKey string `ksql:"node_key"` + InvocationIndex int `ksql:"invocation_index"` + Sequence int64 `ksql:"sequence"` +} + +func workflowAgentTaskFromRow(row *workflowAgentTaskRow) (*WorkflowAgentTaskEntity, error) { + createdAt, err := parseTime(row.CreatedAt) + if err != nil { + return nil, err + } + + return &WorkflowAgentTaskEntity{CreatedAt: createdAt, WorkflowTaskID: row.WorkflowTaskID, + AgentTaskID: row.AgentTaskID, NodeKey: row.NodeKey, + InvocationIndex: row.InvocationIndex, Sequence: row.Sequence}, nil +} diff --git a/internal/database/workflow_repository_test.go b/internal/database/workflow_repository_test.go new file mode 100644 index 00000000..d835fe07 --- /dev/null +++ b/internal/database/workflow_repository_test.go @@ -0,0 +1,504 @@ +package database_test + +import ( + "fmt" + "path/filepath" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/omarluq/librecode/internal/database" +) + +func TestWorkflowRepositoryLifecycleAndAgentLinks(t *testing.T) { + t.Parallel() + + fixture := newTaskTestFixture(t) + ctx := t.Context() + owner, firstChild := fixture.createAgentTaskSessions(ctx) + secondChild, err := fixture.sessions.CreateSession(ctx, owner.CWD, "second", owner.ID) + require.NoError(t, err) + + repository := fixture.workflows + created, err := repository.Create(ctx, newWorkflowRun(owner.ID)) + require.NoError(t, err) + assert.Equal(t, database.TaskKindWorkflow, created.Task.Kind) + assert.Equal(t, database.TaskQueued, created.Task.State) + + loaded, found, err := repository.Get(ctx, created.Task.ID) + require.NoError(t, err) + require.True(t, found) + assert.Equal(t, created.Name, loaded.Name) + assert.Equal(t, created.Source, loaded.Source) + assert.Equal(t, created.ArgumentsJSON, loaded.ArgumentsJSON) + + first, err := fixture.agents.Create(ctx, newAgentTask(owner.ID, firstChild.ID)) + require.NoError(t, err) + second, err := fixture.agents.Create(ctx, newAgentTask(owner.ID, secondChild.ID)) + require.NoError(t, err) + + firstLink, err := repository.LinkAgentTask(ctx, created.Task.ID, first.Task.ID, "inspect", 0) + require.NoError(t, err) + secondLink, err := repository.LinkAgentTask(ctx, created.Task.ID, second.Task.ID, "review", 0) + require.NoError(t, err) + assert.Equal(t, int64(1), firstLink.Sequence) + assert.Equal(t, int64(2), secondLink.Sequence) + + foundLink, found, err := repository.FindAgentTask(ctx, created.Task.ID, " inspect ", 0) + require.NoError(t, err) + require.True(t, found) + assert.Equal(t, first.Task.ID, foundLink.AgentTaskID) + + links, err := repository.ListAgentTasks(ctx, created.Task.ID) + require.NoError(t, err) + require.Len(t, links, 2) + assert.Equal(t, []string{first.Task.ID, second.Task.ID}, []string{links[0].AgentTaskID, links[1].AgentTaskID}) + + changed, err := repository.Tasks().Transition(ctx, created.Task.ID, + []database.TaskState{database.TaskQueued}, database.TaskRunning, taskStartedEvent) + require.NoError(t, err) + require.True(t, changed) + + finish := newTaskFinish( + created.Task.ID, []database.TaskState{database.TaskRunning}, database.TaskSucceeded, taskSucceededEvent, + ) + finish.Result = "complete" + changed, err = repository.Tasks().Finish(ctx, &finish) + require.NoError(t, err) + require.True(t, changed) + + runs, err := repository.ListByOwner(ctx, owner.ID, 10) + require.NoError(t, err) + require.Len(t, runs, 1) + assert.Equal(t, "complete", runs[0].Task.Result) + events, err := repository.Tasks().ListEvents(ctx, created.Task.ID, 0, 10) + require.NoError(t, err) + assert.Len(t, events, 3) +} + +func TestWorkflowRepositoryCreateAgentTaskAtomicallyLinks(t *testing.T) { + t.Parallel() + + fixture := newTaskTestFixture(t) + ctx := t.Context() + owner, child := fixture.createAgentTaskSessions(ctx) + run, err := fixture.workflows.Create(ctx, newWorkflowRun(owner.ID)) + require.NoError(t, err) + + agentTask := newAgentTask(owner.ID, child.ID) + agentTask.Task.ParentTaskID = run.Task.ID + created, err := fixture.workflows.CreateAgentTask(ctx, run.Task.ID, agentTask, "inspect", 0) + require.NoError(t, err) + + link, found, err := fixture.workflows.FindAgentTask(ctx, run.Task.ID, "inspect", 0) + require.NoError(t, err) + require.True(t, found) + assert.Equal(t, created.Task.ID, link.AgentTaskID) + + invalid := newAgentTask(owner.ID, child.ID) + invalid.Task.ParentTaskID = run.Task.ID + _, err = fixture.workflows.CreateAgentTask(ctx, run.Task.ID, invalid, "inspect", 0) + require.Error(t, err) + + listed, listErr := fixture.agents.ListByOwner(ctx, owner.ID, 10) + require.NoError(t, listErr) + assert.Len(t, listed, 1, "failed workflow links must roll back agent task creation") +} + +func TestWorkflowRepositoryCreatesChildSessionAtomically(t *testing.T) { + t.Parallel() + + fixture := newTaskTestFixture(t) + ctx := t.Context() + owner := fixture.createOwner(ctx) + run, err := fixture.workflows.Create(ctx, newWorkflowRun(owner.ID)) + require.NoError(t, err) + + create := func(nodeKey string) (*database.AgentTaskEntity, error) { + task := newAgentTask(owner.ID, "") + task.Task.ParentTaskID = run.Task.ID + + return fixture.workflows.CreateAgentTaskWithChildSession( + ctx, + run.Task.ID, + task, + &database.ChildSessionRequest{CWD: owner.CWD, Name: nodeKey, ParentSessionID: owner.ID}, + "inspect", + 0, + ) + } + + created, err := create("child") + require.NoError(t, err) + assert.NotEmpty(t, created.ChildSessionID) + + _, err = create("orphan") + require.Error(t, err) + + children, err := fixture.sessions.ListChildSessions(ctx, owner.ID) + require.NoError(t, err) + assert.Len(t, children, 1, "link conflicts must roll back their child session") + + agentTasks, err := fixture.agents.ListByOwner(ctx, owner.ID, 10) + require.NoError(t, err) + assert.Len(t, agentTasks, 1, "link conflicts must roll back their agent task") + + links, err := fixture.workflows.ListAgentTasks(ctx, run.Task.ID) + require.NoError(t, err) + assert.Len(t, links, 1, "link conflicts must not add another workflow link") + + mismatched := newAgentTask(owner.ID, "") + mismatched.Task.ParentTaskID = run.Task.ID + _, err = fixture.workflows.CreateAgentTaskWithChildSession( + ctx, run.Task.ID, mismatched, + &database.ChildSessionRequest{CWD: owner.CWD, Name: "wrong-parent", ParentSessionID: testUUIDV7(t)}, + "other", 0, + ) + require.ErrorContains(t, err, "parent differs from agent task owner") +} + +func TestWorkflowRepositoryValidation(t *testing.T) { + t.Parallel() + + const ( + testSource = "source" + testHash = "hash" + argumentsObjectMessage = "arguments_json must be a JSON object" + ) + + fixture := newTaskTestFixture(t) + owner := fixture.createOwner(t.Context()) + repository := fixture.workflows + + tests := []struct { + name string + run *database.WorkflowRunEntity + wantError string + }{ + {name: "nil", run: nil, wantError: "workflow run is required"}, + {name: "source", run: workflowRunForValidation(owner.ID, "", testHash, "{}"), + wantError: "source is required"}, + {name: "hash", run: workflowRunForValidation(owner.ID, testSource, "", "{}"), + wantError: "source_hash is required"}, + {name: "arguments syntax", run: workflowRunForValidation(owner.ID, testSource, testHash, "{"), + wantError: argumentsObjectMessage}, + {name: "arguments array", run: workflowRunForValidation(owner.ID, testSource, testHash, "[]"), + wantError: argumentsObjectMessage}, + {name: "arguments null", run: workflowRunForValidation(owner.ID, testSource, testHash, "null"), + wantError: argumentsObjectMessage}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + _, err := repository.Create(t.Context(), test.run) + require.ErrorContains(t, err, test.wantError) + }) + } + + _, err := repository.LinkAgentTask(t.Context(), "bad", "bad", "", 0) + require.ErrorContains(t, err, "workflow_task_id must be a UUIDv7") + + _, err = repository.CreateAgentTask(t.Context(), testUUIDV7(t), nil, "node", 0) + require.ErrorContains(t, err, "agent task is required") +} + +func TestWorkflowRepositoryLinkAgentTaskOnlyRecoversInvocationUniqueness(t *testing.T) { + t.Parallel() + + ctx := t.Context() + connection := openTestSQLite(t, filepath.Join(t.TempDir(), "workflow.db"), 0) + require.NoError(t, database.Migrate(ctx, connection)) + sessions := database.NewSessionRepository(connection) + workflows := database.NewWorkflowRepository(connection) + agents := database.NewAgentTaskRepository(connection) + + owner, err := sessions.CreateSession(ctx, t.TempDir(), "owner", "") + require.NoError(t, err) + child, err := sessions.CreateSession(ctx, owner.CWD, "child", owner.ID) + require.NoError(t, err) + run, err := workflows.Create(ctx, newWorkflowRun(owner.ID)) + require.NoError(t, err) + task, err := agents.Create(ctx, newAgentTask(owner.ID, child.ID)) + require.NoError(t, err) + _, err = workflows.LinkAgentTask(ctx, run.Task.ID, task.Task.ID, "node", 0) + require.NoError(t, err) + + _, err = connection.ExecContext(ctx, `CREATE TRIGGER reject_workflow_link BEFORE INSERT ON workflow_agent_tasks +BEGIN SELECT RAISE(ABORT, 'reject workflow link'); END`) + require.NoError(t, err) + + _, err = workflows.LinkAgentTask(ctx, run.Task.ID, task.Task.ID, "node", 0) + require.ErrorContains(t, err, "reject workflow link") +} + +func TestWorkflowRepositoryLinkAgentTaskAssignsConcurrentUniqueSequences(t *testing.T) { + t.Parallel() + + fixture := newTaskTestFixture(t) + ctx := t.Context() + owner := fixture.createOwner(ctx) + run, err := fixture.workflows.Create(ctx, newWorkflowRun(owner.ID)) + require.NoError(t, err) + + const count = 8 + + taskIDs := make([]string, count) + for index := range count { + child, createErr := fixture.sessions.CreateSession(ctx, owner.CWD, "child", owner.ID) + require.NoError(t, createErr) + task, createErr := fixture.agents.Create(ctx, newAgentTask(owner.ID, child.ID)) + require.NoError(t, createErr) + + taskIDs[index] = task.Task.ID + } + + var wait sync.WaitGroup + + errors := make(chan error, count) + for index := range count { + wait.Go(func() { + _, linkErr := fixture.workflows.LinkAgentTask(ctx, run.Task.ID, taskIDs[index], "node", index) + errors <- linkErr + }) + } + + wait.Wait() + close(errors) + + for linkErr := range errors { + require.NoError(t, linkErr) + } + + links, err := fixture.workflows.ListAgentTasks(ctx, run.Task.ID) + require.NoError(t, err) + require.Len(t, links, count) + + for index := range links { + assert.Equal(t, int64(index+1), links[index].Sequence) + } +} + +func TestWorkflowRepositoryBoundariesAndIdempotentLinks(t *testing.T) { + t.Parallel() + + fixture := newTaskTestFixture(t) + ctx := t.Context() + owner, child := fixture.createAgentTaskSessions(ctx) + run := newWorkflowRun(owner.ID) + run.ArgumentsJSON = "" + created, err := fixture.workflows.Create(ctx, run) + require.NoError(t, err) + assert.Equal(t, "{}", created.ArgumentsJSON) + assert.NotNil(t, fixture.workflows.AgentTasks()) + + task, err := fixture.agents.Create(ctx, newAgentTask(owner.ID, child.ID)) + require.NoError(t, err) + first, err := fixture.workflows.LinkAgentTask(ctx, created.Task.ID, task.Task.ID, " node ", 0) + require.NoError(t, err) + repeated, err := fixture.workflows.LinkAgentTask(ctx, created.Task.ID, task.Task.ID, "node", 0) + require.NoError(t, err) + assert.Equal(t, first, repeated) + + otherChild, err := fixture.sessions.CreateSession(ctx, owner.CWD, "other", owner.ID) + require.NoError(t, err) + other, err := fixture.agents.Create(ctx, newAgentTask(owner.ID, otherChild.ID)) + require.NoError(t, err) + _, err = fixture.workflows.LinkAgentTask(ctx, created.Task.ID, other.Task.ID, "node", 0) + require.ErrorContains(t, err, "already linked") + + _, err = fixture.workflows.LinkAgentTask(ctx, created.Task.ID, "bad", "node", 1) + require.ErrorContains(t, err, "agent_task_id must be a UUIDv7") +} + +func TestWorkflowRepositoryCreateRollsBackEveryWriteOnFailure(t *testing.T) { + t.Parallel() + + const workflowRunsTable = "workflow_runs" + + tests := []struct { + name string + trigger string + }{ + { + name: "workflow metadata insert", + trigger: `CREATE TRIGGER reject_workflow BEFORE INSERT ON workflow_runs +BEGIN SELECT RAISE(ABORT, 'reject workflow'); END`, + }, + { + name: "initial event insert", + trigger: `CREATE TRIGGER reject_workflow_event BEFORE INSERT ON task_events +WHEN NEW.kind = 'task_queued' +BEGIN SELECT RAISE(ABORT, 'reject workflow event'); END`, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + ctx := t.Context() + connection := openTestSQLite(t, filepath.Join(t.TempDir(), "workflow.db"), 0) + require.NoError(t, database.Migrate(ctx, connection)) + sessions := database.NewSessionRepository(connection) + owner, err := sessions.CreateSession(ctx, t.TempDir(), "owner", "") + require.NoError(t, err) + _, err = connection.ExecContext(ctx, test.trigger) + require.NoError(t, err) + + _, err = database.NewWorkflowRepository(connection).Create(ctx, newWorkflowRun(owner.ID)) + require.Error(t, err) + + for _, table := range []string{"tasks", workflowRunsTable, "task_events"} { + var count int + + err = connection.QueryRowContext(ctx, "SELECT COUNT(*) FROM "+table).Scan(&count) + require.NoError(t, err) + assert.Zero(t, count, "%s write must be rolled back", table) + } + }) + } +} + +func TestWorkflowRepositoryReportsStorageAndCorruptRowErrors(t *testing.T) { + t.Parallel() + + const workflowRunsTable = "workflow_runs" + + t.Run("corrupt workflow timestamp", func(t *testing.T) { + t.Parallel() + + ctx := t.Context() + connection := openTestSQLite(t, filepath.Join(t.TempDir(), "workflow.db"), 0) + require.NoError(t, database.Migrate(ctx, connection)) + + sessions := database.NewSessionRepository(connection) + owner, err := sessions.CreateSession(ctx, t.TempDir(), "owner", "") + require.NoError(t, err) + + repository := database.NewWorkflowRepository(connection) + run, err := repository.Create(ctx, newWorkflowRun(owner.ID)) + require.NoError(t, err) + _, err = connection.ExecContext(ctx, `UPDATE tasks SET created_at = 'not-time' WHERE id = ?`, run.Task.ID) + require.NoError(t, err) + + _, found, err := repository.Get(ctx, run.Task.ID) + assert.False(t, found) + require.ErrorContains(t, err, "parse timestamp") + }) + + tests := []struct { + call func(*database.WorkflowRepository) error + name string + table string + }{ + {call: func(repository *database.WorkflowRepository) error { + _, _, err := repository.Get(t.Context(), testUUIDV7(t)) + + return fmt.Errorf("get workflow: %w", err) + }, name: "load one", table: workflowRunsTable}, + {call: func(repository *database.WorkflowRepository) error { + _, err := repository.ListByOwner(t.Context(), testUUIDV7(t), 0) + + return fmt.Errorf("list workflows: %w", err) + }, name: "list runs", table: workflowRunsTable}, + {call: func(repository *database.WorkflowRepository) error { + _, err := repository.ListAgentTasks(t.Context(), testUUIDV7(t)) + + return fmt.Errorf("list workflow links: %w", err) + }, name: "list links", table: "workflow_agent_tasks"}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + ctx := t.Context() + connection := openTestSQLite(t, filepath.Join(t.TempDir(), "workflow.db"), 0) + require.NoError(t, database.Migrate(ctx, connection)) + _, err := connection.ExecContext(ctx, "DROP TABLE "+test.table) + require.NoError(t, err) + + require.Error(t, test.call(database.NewWorkflowRepository(connection))) + }) + } +} + +func TestWorkflowRepositoryConcurrentIdenticalLinkIsIdempotent(t *testing.T) { + t.Parallel() + + fixture := newTaskTestFixture(t) + ctx := t.Context() + owner, child := fixture.createAgentTaskSessions(ctx) + run, err := fixture.workflows.Create(ctx, newWorkflowRun(owner.ID)) + require.NoError(t, err) + agent, err := fixture.agents.Create(ctx, newAgentTask(owner.ID, child.ID)) + require.NoError(t, err) + + const count = 8 + + var wait sync.WaitGroup + + errors := make(chan error, count) + for range count { + wait.Go(func() { + _, linkErr := fixture.workflows.LinkAgentTask(ctx, run.Task.ID, agent.Task.ID, " same ", 3) + errors <- linkErr + }) + } + + wait.Wait() + close(errors) + + for linkErr := range errors { + require.NoError(t, linkErr) + } + + links, err := fixture.workflows.ListAgentTasks(ctx, run.Task.ID) + require.NoError(t, err) + require.Len(t, links, 1) + assert.Equal(t, int64(1), links[0].Sequence) + assert.Equal(t, "same", links[0].NodeKey) + assert.Equal(t, 3, links[0].InvocationIndex) +} + +func TestWorkflowRepositoryRejectsAgentParentAndOwnerMismatch(t *testing.T) { + t.Parallel() + + fixture := newTaskTestFixture(t) + ctx := t.Context() + owner, child := fixture.createAgentTaskSessions(ctx) + run, err := fixture.workflows.Create(ctx, newWorkflowRun(owner.ID)) + require.NoError(t, err) + + wrongParent := newAgentTask(owner.ID, child.ID) + wrongParent.Task.ParentTaskID = testUUIDV7(t) + _, err = fixture.workflows.CreateAgentTask(ctx, run.Task.ID, wrongParent, "node", 0) + require.ErrorContains(t, err, "agent task parent") + + otherOwner := fixture.createOwner(ctx) + otherChild, err := fixture.sessions.CreateSession(ctx, otherOwner.CWD, "child", otherOwner.ID) + require.NoError(t, err) + + wrongOwner := newAgentTask(otherOwner.ID, otherChild.ID) + wrongOwner.Task.ParentTaskID = run.Task.ID + _, err = fixture.workflows.CreateAgentTask(ctx, run.Task.ID, wrongOwner, "node", 0) + require.ErrorContains(t, err, "owner differs") +} + +func newWorkflowRun(ownerSessionID string) *database.WorkflowRunEntity { + return &database.WorkflowRunEntity{ + Task: *newTask(ownerSessionID), Name: "inspect database", Source: "return agent('inspect')", + SourceHash: "sha256:first", SourceVersion: "v1", ArgumentsJSON: `{"scope":"database"}`, + } +} + +func workflowRunForValidation(ownerSessionID, source, hash, arguments string) *database.WorkflowRunEntity { + return &database.WorkflowRunEntity{ + Task: *newTask(ownerSessionID), Name: "validation", Source: source, SourceHash: hash, SourceVersion: "", + ArgumentsJSON: arguments, + } +} diff --git a/internal/di/agent_task_service.go b/internal/di/agent_task_service.go index b70ae73f..ea822cdb 100644 --- a/internal/di/agent_task_service.go +++ b/internal/di/agent_task_service.go @@ -33,8 +33,8 @@ func NewAgentTaskService(injector do.Injector) (*AgentTaskService, error) { return nil, serviceError(err, "create agent task runner") } - tasks, err := agenttask.New(context.Background(), agenttask.Options{ - Tasks: databaseService.Tasks, AgentTasks: databaseService.AgentTasks, + tasks, err := agenttask.New(context.Background(), &agenttask.Options{ + Tasks: databaseService.Tasks, AgentTasks: databaseService.AgentTasks, Workflows: databaseService.Workflows, Runner: runner, Concurrency: 0, SessionConcurrency: 0, QueueCapacity: 0, Timeout: 0, Logger: logger, }) diff --git a/internal/di/assistant_service_internal_test.go b/internal/di/assistant_service_internal_test.go index 4139c7c5..7086cad5 100644 --- a/internal/di/assistant_service_internal_test.go +++ b/internal/di/assistant_service_internal_test.go @@ -75,6 +75,7 @@ func newTestDatabaseService(t *testing.T) *DatabaseService { Documents: database.NewDocumentRepository(connection), Tasks: database.NewTaskRepository(connection), AgentTasks: database.NewAgentTaskRepository(connection), + Workflows: database.NewWorkflowRepository(connection), path: "", } } diff --git a/internal/di/chat_workflow_service.go b/internal/di/chat_workflow_service.go new file mode 100644 index 00000000..b5970fb8 --- /dev/null +++ b/internal/di/chat_workflow_service.go @@ -0,0 +1,45 @@ +package di + +import ( + "context" + "log/slog" + + "github.com/samber/do/v2" + + "github.com/omarluq/librecode/internal/workflow" +) + +// ChatWorkflowService owns the in-process workflow dispatcher used by interactive chat. +type ChatWorkflowService struct { + Runs *workflow.Service + Dispatcher *workflow.Dispatcher +} + +// NewChatWorkflowService enables model-authored workflows for interactive chat. +func NewChatWorkflowService(injector do.Injector) (*ChatWorkflowService, error) { + workflows := do.MustInvoke[*WorkflowService](injector) + databaseService := do.MustInvoke[*DatabaseService](injector) + assistantService := do.MustInvoke[*AssistantService](injector) + + var logger *slog.Logger + if loggerService, err := do.Invoke[*LoggerService](injector); err == nil { + logger = loggerService.SlogLogger + } + + dispatcher, err := workflow.NewDispatcher(context.Background(), workflow.DispatcherOptions{ + Service: workflows.Runs, Tasks: databaseService.Tasks, Logger: logger, + Concurrency: 0, Buffer: 0, Interval: 0, + }) + if err != nil { + return nil, serviceError(err, "create chat workflow dispatcher") + } + + assistantService.Runtime.SetWorkflowSubmitter(dispatcher) + + return &ChatWorkflowService{Runs: workflows.Runs, Dispatcher: dispatcher}, nil +} + +// Shutdown stops workflow workers before their dependencies are closed. +func (service *ChatWorkflowService) Shutdown(ctx context.Context) error { + return serviceError(service.Dispatcher.Shutdown(ctx), "shutdown chat workflow dispatcher") +} diff --git a/internal/di/container.go b/internal/di/container.go index 7ee96bc8..11ea9b30 100644 --- a/internal/di/container.go +++ b/internal/di/container.go @@ -73,6 +73,16 @@ func (c *Container) AgentTaskService() *AgentTaskService { return do.MustInvoke[*AgentTaskService](c.injector) } +// WorkflowService resolves script-driven durable agent orchestration. +func (c *Container) WorkflowService() *WorkflowService { + return do.MustInvoke[*WorkflowService](c.injector) +} + +// ChatWorkflowService resolves the interactive workflow dispatcher. +func (c *Container) ChatWorkflowService() *ChatWorkflowService { + return do.MustInvoke[*ChatWorkflowService](c.injector) +} + // ToolService resolves the tool service. func (c *Container) ToolService() *ToolService { return do.MustInvoke[*ToolService](c.injector) diff --git a/internal/di/database_service.go b/internal/di/database_service.go index 7ddba2bd..02e04551 100644 --- a/internal/di/database_service.go +++ b/internal/di/database_service.go @@ -30,6 +30,7 @@ type DatabaseService struct { Documents *database.DocumentRepository Tasks *database.TaskRepository AgentTasks *database.AgentTaskRepository + Workflows *database.WorkflowRepository path string } @@ -63,6 +64,7 @@ func NewDatabaseService(injector do.Injector) (*DatabaseService, error) { Documents: database.NewDocumentRepositoryWithProvider(sqlProvider), Tasks: database.NewTaskRepositoryWithProvider(sqlProvider), AgentTasks: database.NewAgentTaskRepositoryWithProvider(sqlProvider), + Workflows: database.NewWorkflowRepositoryWithProvider(sqlProvider), path: databasePath, }, nil } diff --git a/internal/di/database_service_internal_test.go b/internal/di/database_service_internal_test.go index 2cce32c2..b7ce3a06 100644 --- a/internal/di/database_service_internal_test.go +++ b/internal/di/database_service_internal_test.go @@ -54,7 +54,9 @@ func TestDatabaseServiceHealthCheckAndShutdown(t *testing.T) { connection, err := sql.Open("sqlite", filepath.Join(t.TempDir(), "librecode.db")) require.NoError(t, err) - service := &DatabaseService{DB: connection, Sessions: nil, Documents: nil, Tasks: nil, AgentTasks: nil, path: ""} + service := &DatabaseService{ + DB: connection, Sessions: nil, Documents: nil, Tasks: nil, AgentTasks: nil, Workflows: nil, path: "", + } require.NoError(t, service.HealthCheck(context.Background())) require.NoError(t, service.Shutdown(context.Background())) @@ -68,7 +70,9 @@ func TestDatabaseServiceShutdownReturnsContextError(t *testing.T) { require.NoError(t, err) t.Cleanup(func() { require.NoError(t, connection.Close()) }) - service := &DatabaseService{DB: connection, Sessions: nil, Documents: nil, Tasks: nil, AgentTasks: nil, path: ""} + service := &DatabaseService{ + DB: connection, Sessions: nil, Documents: nil, Tasks: nil, AgentTasks: nil, Workflows: nil, path: "", + } ctx, cancel := context.WithCancel(context.Background()) cancel() diff --git a/internal/di/model_service_internal_test.go b/internal/di/model_service_internal_test.go index 4e4f16f6..23ba5d6b 100644 --- a/internal/di/model_service_internal_test.go +++ b/internal/di/model_service_internal_test.go @@ -38,8 +38,8 @@ func TestNewModelServiceWiresRegistryDiscovery(t *testing.T) { Sessions: nil, Documents: documents, Tasks: nil, - AgentTasks: nil, - path: "", + AgentTasks: nil, Workflows: nil, + path: "", }) do.ProvideValue(injector, &AuthService{Storage: storage}) diff --git a/internal/di/register.go b/internal/di/register.go index c86d9ef6..c4add311 100644 --- a/internal/di/register.go +++ b/internal/di/register.go @@ -15,4 +15,6 @@ func RegisterServices(injector do.Injector) { do.Provide(injector, NewToolService) do.Provide(injector, NewAssistantService) do.Provide(injector, NewAgentTaskService) + do.Provide(injector, NewWorkflowService) + do.Provide(injector, NewChatWorkflowService) } diff --git a/internal/di/service_constructors_internal_test.go b/internal/di/service_constructors_internal_test.go index 9ee2eaed..1f74f05b 100644 --- a/internal/di/service_constructors_internal_test.go +++ b/internal/di/service_constructors_internal_test.go @@ -51,7 +51,7 @@ func TestNewAgentTaskServiceRejectsIncompleteWiring(t *testing.T) { injector := do.New() do.ProvideValue(injector, &DatabaseService{ - DB: nil, Sessions: nil, Documents: nil, Tasks: nil, AgentTasks: nil, path: "", + DB: nil, Sessions: nil, Documents: nil, Tasks: nil, AgentTasks: nil, Workflows: nil, path: "", }) do.ProvideValue(injector, &AssistantService{Runtime: nil, Agents: nil}) @@ -115,6 +115,8 @@ func TestContainerServiceAccessors(t *testing.T) { require.NotNil(t, container.ModelService()) require.NotNil(t, container.AssistantService()) require.NotNil(t, container.AgentTaskService()) + require.NotNil(t, container.WorkflowService()) + require.NotNil(t, container.ChatWorkflowService()) require.NotNil(t, container.ToolService()) require.NotNil(t, container.SkillsService()) } diff --git a/internal/di/workflow_service.go b/internal/di/workflow_service.go new file mode 100644 index 00000000..09525feb --- /dev/null +++ b/internal/di/workflow_service.go @@ -0,0 +1,56 @@ +package di + +import ( + "context" + + "github.com/samber/do/v2" + + "github.com/omarluq/librecode/internal/assistant" + "github.com/omarluq/librecode/internal/workflow" +) + +// WorkflowService owns script-driven durable agent orchestration. +type WorkflowService struct { + Runner *workflow.Runner + Runs *workflow.Service +} + +// NewWorkflowService wires workflow execution to the durable agent scheduler. +func NewWorkflowService(injector do.Injector) (*WorkflowService, error) { + databaseService := do.MustInvoke[*DatabaseService](injector) + assistantService := do.MustInvoke[*AssistantService](injector) + agentTaskService := do.MustInvoke[*AgentTaskService](injector) + + submitter, err := assistant.NewAgentSubmitter( + agentTaskService.Tasks, + assistantService.Agents, + ) + if err != nil { + return nil, serviceError(err, "create workflow agent submitter") + } + + controller, err := assistant.NewWorkflowController( + submitter, + agentTaskService.Tasks, + databaseService.Sessions, + ) + if err != nil { + return nil, serviceError(err, "create workflow controller") + } + + runner, err := workflow.NewRunner(controller) + if err != nil { + return nil, serviceError(err, "create workflow runner") + } + + runs, err := workflow.NewService(databaseService.Workflows, runner) + if err != nil { + return nil, serviceError(err, "create durable workflow service") + } + + if _, recoverErr := runs.RecoverInterrupted(context.Background()); recoverErr != nil { + return nil, serviceError(recoverErr, "recover interrupted workflows") + } + + return &WorkflowService{Runner: runner, Runs: runs}, nil +} diff --git a/internal/executeworker/client.go b/internal/executeworker/client.go new file mode 100644 index 00000000..ed8c7585 --- /dev/null +++ b/internal/executeworker/client.go @@ -0,0 +1,328 @@ +package executeworker + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "sync" + + "github.com/omarluq/librecode/internal/mvmhost" +) + +// RPCHandler handles a callback request received from an execute worker. +type RPCHandler func(context.Context, *Message) (any, error) + +// Client evaluates source in a separate worker process. +type Client struct { + Handler RPCHandler + Executable string +} + +type workerProcess struct { + killErr error + cmd *exec.Cmd + killOnce sync.Once +} + +// Request describes one isolated MVM evaluation. +type Request struct { + Arguments any + Mode string + Name string + Source string +} + +// Eval evaluates provider-facing source, forwarding callback requests to Handler. +func (client Client) Eval(ctx context.Context, source string) (mvmhost.Result, error) { + return client.EvalRequest(ctx, &Request{ + Arguments: nil, Mode: "execute", Name: "execute.go", Source: source, + }) +} + +// EvalRequest evaluates source in the requested worker mode. +func (client Client) EvalRequest(ctx context.Context, eval *Request) (mvmhost.Result, error) { + worker, stdin, stdout, err := client.startWorker() + if err != nil { + return mvmhost.Result{}, err + } + + stopCancellation := make(chan struct{}) + + go func() { + select { + case <-ctx.Done(): + worker.kill() + case <-stopCancellation: + } + }() + + defer close(stopCancellation) + + if eval == nil { + return mvmhost.Result{}, worker.abort(errors.New("execute worker request is required")) + } + + request := newMessage("eval") + request.Mode, request.Name, request.Source = eval.Mode, eval.Name, eval.Source + + if request.Arguments, err = json.Marshal(eval.Arguments); err != nil { + return mvmhost.Result{}, worker.abort(fmt.Errorf("encode worker arguments: %w", err)) + } + + if err = Write(stdin, &request); err != nil { + return mvmhost.Result{}, worker.abort(err) + } + + return client.readMessages(ctx, worker, stdin, stdout) +} + +func (client Client) startWorker() (*workerProcess, io.WriteCloser, io.ReadCloser, error) { + executable, err := client.executablePath() + if err != nil { + return nil, nil, nil, err + } + + // Constructing Cmd directly avoids shell interpretation. executablePath only + // permits an absolute path to a regular file. + cmd := &exec.Cmd{Path: executable, Args: []string{executable, "__execute-worker"}, Env: os.Environ()} + + stdin, err := cmd.StdinPipe() + if err != nil { + return nil, nil, nil, fmt.Errorf("open execute worker stdin: %w", err) + } + + stdout, err := cmd.StdoutPipe() + if err != nil { + return nil, nil, nil, errors.Join( + fmt.Errorf("open execute worker stdout: %w", err), + stdin.Close(), + ) + } + + cmd.Stderr = os.Stderr + if err = cmd.Start(); err != nil { + return nil, nil, nil, errors.Join( + fmt.Errorf("start execute worker: %w", err), + stdin.Close(), + stdout.Close(), + ) + } + + return &workerProcess{cmd: cmd, killOnce: sync.Once{}, killErr: nil}, stdin, stdout, nil +} + +func (client Client) executablePath() (string, error) { + executable := client.Executable + if executable == "" { + var err error + + executable, err = os.Executable() + if err != nil { + return "", fmt.Errorf("resolve execute worker executable: %w", err) + } + } + + if !filepath.IsAbs(executable) || filepath.Clean(executable) != executable { + return "", fmt.Errorf("execute worker path must be an absolute clean path: %q", executable) + } + + info, err := os.Stat(executable) + if err != nil { + return "", fmt.Errorf("inspect execute worker executable: %w", err) + } + + if !info.Mode().IsRegular() { + return "", fmt.Errorf("execute worker path is not a regular file: %q", executable) + } + + return executable, nil +} + +func (client Client) readMessages( + ctx context.Context, + worker *workerProcess, + stdin io.WriteCloser, + stdout io.Reader, +) (mvmhost.Result, error) { + var ( + writes sync.Mutex + callbacks []<-chan struct{} + ) + + for { + message, err := Read(stdout) + if err != nil { + return mvmhost.Result{}, worker.readError(ctx, err) + } + + switch message.Type { + case "rpc": + callbackDone := make(chan struct{}) + + callbacks = append(callbacks, callbackDone) + go func(rpc Message) { + defer close(callbackDone) + + response := client.rpcResponse(ctx, &rpc) + + writes.Lock() + defer writes.Unlock() + + if Write(stdin, &response) != nil { + worker.kill() + } + }(message) + case "result": + if err := waitForCallbacks(ctx, callbacks); err != nil { + return mvmhost.Result{}, worker.abort(err) + } + + return finishResult(ctx, worker, stdin, &message) + default: + return mvmhost.Result{}, worker.abort( + fmt.Errorf("unexpected execute worker message %q", message.Type), + ) + } + } +} + +func waitForCallbacks(ctx context.Context, callbacks []<-chan struct{}) error { + for _, callbackDone := range callbacks { + select { + case <-callbackDone: + case <-ctx.Done(): + return canceledError(ctx.Err()) + } + } + + return nil +} + +func (client Client) rpcResponse(ctx context.Context, message *Message) Message { + var ( + value any + rpcErr error + ) + if client.Handler == nil { + rpcErr = errors.New("execute worker RPC handler is not configured") + } else { + value, rpcErr = client.Handler(ctx, message) + } + + response := newMessage("rpc_result") + + response.ID = message.ID + if rpcErr != nil { + response.Error = rpcErr.Error() + } + + if _, ok := value.(ToolCallResult); ok { + response.ValueKind = toolCallResultKind + } + + if response.Value, rpcErr = json.Marshal(value); rpcErr != nil && response.Error == "" { + response.Error = rpcErr.Error() + } + + return response +} + +func finishResult(ctx context.Context, + worker *workerProcess, stdin io.Closer, message *Message) (mvmhost.Result, error) { + if err := stdin.Close(); err != nil { + return mvmhost.Result{}, worker.abort(fmt.Errorf("close execute worker stdin: %w", err)) + } + + if err := waitForWorker(ctx, worker); err != nil { + return mvmhost.Result{}, err + } + + result := mvmhost.Result{ + Value: nil, ValueKind: message.ValueKind, Stdout: message.Stdout, Stderr: message.Stderr, + } + if len(message.Value) > 0 { + switch { + case string(message.Value) == jsonNullValue: + result.Value = json.RawMessage(jsonNullValue) + case message.ValueKind == toolCallResultKind: + var nested ToolCallResult + if err := json.Unmarshal(message.Value, &nested); err != nil { + return result, fmt.Errorf("decode execute worker tool result: %w", err) + } + + result.Value = nested + default: + if err := json.Unmarshal(message.Value, &result.Value); err != nil { + return result, fmt.Errorf("decode execute worker result: %w", err) + } + } + } + + if message.Error == "" { + return result, nil + } + + return result, &mvmhost.EvalError{ + Err: errors.New(message.Error), Kind: mvmhost.ErrorKind(message.ErrorKind), ExitCode: message.ExitCode, + } +} + +func waitForWorker(ctx context.Context, worker *workerProcess) error { + if err := worker.cmd.Wait(); err != nil { + if ctx.Err() != nil { + return canceledError(ctx.Err()) + } + + return fmt.Errorf("wait for execute worker: %w", err) + } + + return nil +} + +func (worker *workerProcess) kill() { + worker.killOnce.Do(func() { + worker.killErr = worker.cmd.Process.Kill() + }) +} + +func (worker *workerProcess) abort(cause error) error { + worker.kill() + + waitErr := worker.cmd.Wait() + if worker.killErr != nil && !errors.Is(worker.killErr, os.ErrProcessDone) { + return errors.Join(cause, fmt.Errorf("kill execute worker: %w", worker.killErr), waitErr) + } + + return errors.Join(cause, waitErr) +} + +func (worker *workerProcess) readError(ctx context.Context, readErr error) error { + if !errors.Is(readErr, io.EOF) || ctx.Err() != nil { + worker.kill() + } + + waitErr := worker.cmd.Wait() + + if ctx.Err() != nil { + return canceledError(ctx.Err()) + } + + if !errors.Is(readErr, io.EOF) { + return errors.Join(fmt.Errorf("read execute worker: %w", readErr), waitErr) + } + + if waitErr != nil { + return fmt.Errorf("execute worker exited without result: %w", waitErr) + } + + return errors.New("execute worker exited without result") +} + +func canceledError(err error) error { + return &mvmhost.EvalError{Err: err, Kind: mvmhost.ErrorKindCanceled, ExitCode: 0} +} diff --git a/internal/executeworker/client_test.go b/internal/executeworker/client_test.go new file mode 100644 index 00000000..3bf6d3bd --- /dev/null +++ b/internal/executeworker/client_test.go @@ -0,0 +1,101 @@ +package executeworker_test + +import ( + "context" + "encoding/json" + "errors" + "os" + "testing" + "time" + + "github.com/omarluq/librecode/internal/executeworker" + "github.com/omarluq/librecode/internal/tool" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const ( + helperEnv = "LIBRECODE_EXECUTEWORKER_TEST_HELPER" + echoQuery = "echo" +) + +func TestMain(testMain *testing.M) { + if os.Getenv(helperEnv) == "1" { + if err := executeworker.Serve(os.Stdin, os.Stdout); err != nil { + os.Exit(2) + } + + os.Exit(0) + } + + os.Exit(testMain.Run()) +} + +func testClient() executeworker.Client { + return executeworker.Client{Executable: os.Args[0], Handler: nil} +} + +func TestClientHardCancelsInfiniteLoop(t *testing.T) { + t.Setenv(helperEnv, "1") + + ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond) + defer cancel() + + started := time.Now() + _, err := testClient().Eval(ctx, `for {}; 1`) + require.Error(t, err) + require.ErrorIs(t, err, context.DeadlineExceeded) + assert.Less(t, time.Since(started), 3*time.Second) +} + +func TestClientPreservesTypedToolCallResults(t *testing.T) { + t.Setenv(helperEnv, "1") + + want := executeworker.ToolCallResult{ + Details: map[string]any{"path": "image.png"}, + Content: []tool.ContentBlock{ + {Type: tool.ContentTypeText, Text: "caption", Data: "", MIMEType: ""}, + {Type: tool.ContentTypeImage, Text: "", Data: "aW1hZ2U=", MIMEType: "image/png"}, + }, + Error: "nested failure", IsError: true, + } + client := testClient() + client.Handler = func(_ context.Context, _ *executeworker.Message) (any, error) { return want, nil } + + result, err := client.Eval(t.Context(), `import "tools"; tools.Call("image", nil)`) + require.NoError(t, err) + assert.Equal(t, want, result.Value) +} + +func TestClientPreservesNullCallbackResult(t *testing.T) { + t.Setenv(helperEnv, "1") + + client := testClient() + client.Handler = func(_ context.Context, _ *executeworker.Message) (any, error) { + return json.RawMessage("null"), nil + } + + result, err := client.Eval(t.Context(), `import "tools"; tools.Describe("missing")`) + require.NoError(t, err) + + raw, ok := result.Value.(json.RawMessage) + require.True(t, ok) + assert.JSONEq(t, "null", string(raw)) +} + +func TestClientCallbackRPC(t *testing.T) { + t.Setenv(helperEnv, "1") + + client := testClient() + client.Handler = func(_ context.Context, message *executeworker.Message) (any, error) { + if message.Method != "search" || message.Query != echoQuery { + return nil, errors.New("unexpected RPC") + } + + return []map[string]any{{"name": echoQuery}}, nil + } + + result, err := client.Eval(t.Context(), `import "tools"; tools.Search("echo")`) + require.NoError(t, err) + assert.Equal(t, []any{map[string]any{"name": echoQuery}}, result.Value) +} diff --git a/internal/executeworker/coverage_internal_test.go b/internal/executeworker/coverage_internal_test.go new file mode 100644 index 00000000..24b76436 --- /dev/null +++ b/internal/executeworker/coverage_internal_test.go @@ -0,0 +1,532 @@ +package executeworker + +import ( + "bytes" + "context" + "encoding/binary" + "encoding/json" + "errors" + "io" + "os" + "os/exec" + "path/filepath" + "strings" + "sync" + "testing" + + "github.com/omarluq/librecode/internal/mvmhost" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type errorWriter struct{ calls int } + +func (w *errorWriter) Write(p []byte) (int, error) { + w.calls++ + if w.calls == 1 { + return len(p), nil + } + + return 0, errors.New("write failed") +} + +type shortWriter struct{} + +func (shortWriter) Write([]byte) (int, error) { return 0, nil } + +type closeError struct{ err error } + +func (c closeError) Close() error { return c.err } + +type nopWriteCloser struct{ io.Writer } + +func (nopWriteCloser) Close() error { return nil } + +func TestProtocolErrorBranches(t *testing.T) { + t.Parallel() + + _, err := Read(strings.NewReader("x")) + require.ErrorContains(t, err, "frame size") + + for _, size := range []uint32{0, MaxFrameSize + 1} { + var input bytes.Buffer + require.NoError(t, binary.Write(&input, binary.BigEndian, size)) + _, err = Read(&input) + require.ErrorContains(t, err, "exceeds limit") + } + + var truncated bytes.Buffer + require.NoError(t, binary.Write(&truncated, binary.BigEndian, uint32(2))) + require.NoError(t, truncated.WriteByte('{')) + _, err = Read(&truncated) + require.ErrorContains(t, err, "read execute worker frame") + + for _, payload := range []struct { + value string + size uint32 + }{ + {value: "{", size: 1}, + {value: `{`, size: 1}, + } { + var input bytes.Buffer + require.NoError(t, binary.Write(&input, binary.BigEndian, payload.size)) + _, writeErr := input.WriteString(payload.value) + require.NoError(t, writeErr) + + _, err = Read(&input) + require.Error(t, err) + } + + require.Error(t, Write(io.Discard, nil)) + require.Error(t, Write(io.Discard, ptrMessage(""))) + require.ErrorContains(t, Write(shortWriter{}, ptrMessage("rpc")), "frame size") + require.ErrorContains(t, Write(&errorWriter{calls: 0}, ptrMessage("rpc")), "frame") + + var partial bytes.Buffer + + writer := struct{ io.Writer }{Writer: &partial} + require.NoError(t, Write(writer, ptrMessage("rpc"))) + + _, err = Read(&partial) + require.NoError(t, err) +} + +func ptrMessage(kind string) *Message { + message := newMessage(kind) + + return &message +} + +func TestWorkerPipelineBranches(t *testing.T) { + t.Parallel() + + _, err := workerPipeline(nil, func(any) (any, error) { return 0, nil }, 0) + require.EqualError(t, err, "pipeline concurrency must be positive") + _, err = workerPipeline(nil, nil, 1) + require.EqualError(t, err, "pipeline callback is required") + + results, err := workerPipeline([]any{1, 2, 3}, func(value any) (any, error) { + if value == 2 { + return nil, errors.New("stop") + } + + integer, integerOK := value.(int) + require.True(t, integerOK) + + return integer * 2, nil + }, 1) + require.NoError(t, err) + assert.Equal(t, 2, results[0]["value"]) + assert.Equal(t, "stop", results[1]["error"]) + assert.Contains(t, results[2]["error"], "stopped") + + empty, err := workerPipeline(nil, func(any) (any, error) { return 0, nil }, 2) + require.NoError(t, err) + assert.Empty(t, empty) +} + +func TestWorkerBindingsModes(t *testing.T) { + t.Parallel() + + caller := newRPCCaller() + bindings, err := workerBindings(ptrMessage("eval"), caller) + require.NoError(t, err) + assert.Contains(t, bindings, "tools") + + request := ptrMessage("eval") + request.Mode = "bad" + _, err = workerBindings(request, caller) + require.ErrorContains(t, err, "unknown") + + request.Mode = "workflow" + request.Arguments = json.RawMessage(`{`) + _, err = workerBindings(request, caller) + require.ErrorContains(t, err, "decode workflow arguments") + + for _, arguments := range []json.RawMessage{nil, json.RawMessage("null"), json.RawMessage(`{"x":1}`)} { + request.Arguments = arguments + bindings, err = workerBindings(request, caller) + require.NoError(t, err) + assert.Contains(t, bindings, "librecode/workflow") + } +} + +func TestRPCCallerExchangeAndFailures(t *testing.T) { + t.Parallel() + + parentRead, workerWrite := io.Pipe() + workerRead, parentWrite := io.Pipe() + caller := newRPCCaller() + + caller.in, caller.out = workerRead, workerWrite + go caller.readResponses() + + writeErrors := make(chan error, 1) + + go func() { + request, readErr := Read(parentRead) + if readErr != nil { + return + } + + ignored := newMessage("noise") + if writeErr := Write(parentWrite, &ignored); writeErr != nil { + writeErrors <- writeErr + + return + } + + response := newMessage("rpc_result") + response.ID = request.ID + response.Value = json.RawMessage(`{"ok":true}`) + + writeErrors <- Write(parentWrite, &response) + }() + + value, err := caller.callResult("call", "name", "query", map[string]any{"x": 1}) + require.NoError(t, err) + require.NoError(t, <-writeErrors) + assert.Equal(t, map[string]any{"ok": true}, value) + require.NoError(t, parentWrite.Close()) + require.NoError(t, parentRead.Close()) + require.NoError(t, workerWrite.Close()) + + badInput := newRPCCaller() + _, err = badInput.callResult("call", "", "", func() {}) + require.ErrorContains(t, err, "encode worker RPC input") + + terminal := errors.New("terminal") + failed := newRPCCaller() + failed.terminalErr = terminal + _, err = failed.exchange("call", "", "", nil) + require.ErrorIs(t, err, terminal) + + writeFailed := newRPCCaller() + writeFailed.out = shortWriter{} + _, err = writeFailed.exchange("call", "", "", nil) + require.Error(t, err) + assert.Empty(t, writeFailed.pending) + + responseCh := make(chan Message, 1) + pending := newRPCCaller() + pending.pending[1] = responseCh + pending.failPending(errors.New("first")) + pending.failPending(errors.New("second")) + require.EqualError(t, pending.terminalErr, "first") + assert.Equal(t, "first", (<-responseCh).Error) +} + +func TestRPCValuesAndResultMessages(t *testing.T) { + t.Parallel() + + value, err := decodeRPCValue(messageWithValue("", json.RawMessage("null"))) + require.NoError(t, err) + + rawValue, rawValueOK := value.(json.RawMessage) + require.True(t, rawValueOK) + assert.JSONEq(t, "null", string(rawValue)) + + toolValue := ToolCallResult{Details: nil, Error: "nested", Content: nil, IsError: true} + raw, err := json.Marshal(toolValue) + require.NoError(t, err) + value, err = decodeRPCValue(messageWithValue(toolCallResultKind, raw)) + require.NoError(t, err) + assert.Equal(t, toolValue, value) + + _, err = decodeRPCValue(messageWithValue(toolCallResultKind, json.RawMessage("{"))) + require.ErrorContains(t, err, "decode tool call result") + _, err = decodeRPCValue(messageWithValue("", json.RawMessage("{"))) + require.ErrorContains(t, err, "decode RPC result") + + isError, isErrorOK := rpcError("bad")["is_error"].(bool) + require.True(t, isErrorOK) + assert.True(t, isError) + + message := resultMessage(mvmhost.Result{ + Value: pipelineValue{{"x": 1}}, ValueKind: "", Stdout: "out", Stderr: "err", + }, nil) + assert.Equal(t, pipelineResultKind, message.ValueKind) + assert.NotEmpty(t, message.Value) + message = resultMessage(mvmResult(toolValue), nil) + assert.Equal(t, toolCallResultKind, message.ValueKind) + message = resultMessage(mvmResult(func() {}), nil) + assert.Equal(t, string(mvmhost.ErrorKindRuntime), message.ErrorKind) + + evalErr := &mvmhost.EvalError{Err: errors.New("boom"), Kind: mvmhost.ErrorKindCanceled, ExitCode: 7} + message = resultMessage(mvmResult(nil), evalErr) + assert.Equal(t, 7, message.ExitCode) + message = resultMessage(mvmResult(nil), errors.New("plain")) + assert.Equal(t, "plain", message.Error) +} + +func TestServeBranchesInProcess(t *testing.T) { + t.Parallel() + + require.Error(t, Serve(strings.NewReader("bad"), io.Discard)) + + wrongMode := newMessage("eval") + wrongMode.Mode = "wrong" + + for _, request := range []Message{newMessage("wrong"), wrongMode} { + var input bytes.Buffer + require.NoError(t, Write(&input, &request)) + require.Error(t, Serve(&input, io.Discard)) + } + + request := newMessage("eval") + request.Source = "1" + + var input, output bytes.Buffer + require.NoError(t, Write(&input, &request)) + require.NoError(t, Serve(&input, &output)) + response, err := Read(&output) + require.NoError(t, err) + assert.Equal(t, "result", response.Type) +} + +func startedProcess(t *testing.T, script string) *workerProcess { + t.Helper() + + var cmd *exec.Cmd + + switch script { + case "exit 0": + cmd = exec.CommandContext(t.Context(), "sh", "-c", "exit 0") + case "exit 2": + cmd = exec.CommandContext(t.Context(), "sh", "-c", "exit 2") + case "exit 4": + cmd = exec.CommandContext(t.Context(), "sh", "-c", "exit 4") + case "sleep 10": + cmd = exec.CommandContext(t.Context(), "sh", "-c", "sleep 10") + default: + require.FailNow(t, "unsupported test process", script) + } + + require.NoError(t, cmd.Start()) + + return &workerProcess{killErr: nil, cmd: cmd, killOnce: sync.Once{}} +} + +func TestEvalRequestValidationInProcess(t *testing.T) { + t.Setenv("LIBRECODE_EXECUTEWORKER_TEST_HELPER", "1") + + client := newClient(nil, os.Args[0]) + + _, err := client.EvalRequest(t.Context(), nil) + require.ErrorContains(t, err, "request is required") + _, err = client.EvalRequest(t.Context(), requestWithArguments(func() {})) + require.ErrorContains(t, err, "encode worker arguments") +} + +func TestClientPrivateBranches(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + relative := newClient(nil, "relative") + _, err := relative.executablePath() + require.ErrorContains(t, err, "absolute clean path") + _, err = (newClient(nil, filepath.Join(dir, "missing"))).executablePath() + require.ErrorContains(t, err, "inspect") + _, err = (newClient(nil, dir)).executablePath() + require.ErrorContains(t, err, "not a regular file") + + file := filepath.Join(dir, "worker") + require.NoError(t, os.WriteFile(file, []byte("x"), 0o600)) + path, err := (newClient(nil, file)).executablePath() + require.NoError(t, err) + assert.Equal(t, file, path) + path, err = newClient(nil, "").executablePath() + require.NoError(t, err) + assert.True(t, filepath.IsAbs(path)) + + response := newClient(nil, "").rpcResponse(t.Context(), messageWithID(2)) + assert.Contains(t, response.Error, "not configured") + + client := newClient(func(context.Context, *Message) (any, error) { return func() {}, nil }, "") + response = client.rpcResponse(t.Context(), ptrMessage("")) + assert.NotEmpty(t, response.Error) + + client.Handler = func(context.Context, *Message) (any, error) { return nil, errors.New("handler") } + response = client.rpcResponse(t.Context(), ptrMessage("")) + assert.Equal(t, "handler", response.Error) + + unexpected := newMessage("unexpected") + + var output bytes.Buffer + require.NoError(t, Write(&output, &unexpected)) + worker := startedProcess(t, "sleep 10") + _, err = newClient(nil, "").readMessages(t.Context(), worker, nopWriteCloser{io.Discard}, &output) + require.ErrorContains(t, err, "unexpected execute worker message") +} + +func TestReadMessagesCallbackWaitHonorsCancellation(t *testing.T) { + t.Parallel() + + callbackStarted := make(chan struct{}) + releaseCallback := make(chan struct{}) + client := newClient(func(context.Context, *Message) (any, error) { + close(callbackStarted) + <-releaseCallback + + return "released", nil + }, "") + + var output bytes.Buffer + + rpc := newMessage("rpc") + rpc.ID = 1 + require.NoError(t, Write(&output, &rpc)) + + result := newMessage("result") + require.NoError(t, Write(&output, &result)) + + ctx, cancel := context.WithCancel(t.Context()) + worker := startedProcess(t, "sleep 10") + resultCh := make(chan error, 1) + + go func() { + _, err := client.readMessages(ctx, worker, nopWriteCloser{io.Discard}, &output) + resultCh <- err + }() + + <-callbackStarted + cancel() + + select { + case err := <-resultCh: + require.ErrorIs(t, err, context.Canceled) + case <-t.Context().Done(): + require.FailNow(t, "callback wait did not honor cancellation") + } + + close(releaseCallback) +} + +func TestFinishResultBranches(t *testing.T) { + t.Parallel() + + worker := startedProcess(t, "sleep 10") + _, err := finishResult(t.Context(), worker, closeError{err: errors.New("close")}, ptrMessage("result")) + require.ErrorContains(t, err, "close execute worker stdin") + + cases := []struct { + check func(mvmhost.Result, error) + message Message + }{ + {message: *messageWithValue("", json.RawMessage("null")), check: func(result mvmhost.Result, err error) { + require.NoError(t, err) + + rawResult, rawResultOK := result.Value.(json.RawMessage) + require.True(t, rawResultOK) + assert.JSONEq(t, "null", string(rawResult)) + }}, + {message: *messageWithValue("", json.RawMessage(`{"x":1}`)), check: func(result mvmhost.Result, err error) { + require.NoError(t, err) + assert.Equal(t, map[string]any{"x": float64(1)}, result.Value) + }}, + { + message: *messageWithValue( + toolCallResultKind, + json.RawMessage(`{"details":{},"content":[],"is_error":false}`), + ), + check: func(result mvmhost.Result, err error) { + require.NoError(t, err) + assert.IsType(t, emptyToolCallResult(), result.Value) + }, + }, + {message: *messageWithValue("", json.RawMessage("{")), check: func(_ mvmhost.Result, err error) { + require.ErrorContains(t, err, "decode execute worker result") + }}, + { + message: *messageWithValue(toolCallResultKind, json.RawMessage("{")), + check: func(_ mvmhost.Result, err error) { + require.ErrorContains(t, err, "decode execute worker tool result") + }, + }, + {message: *errorMessage("boom", "runtime", 3), check: func(_ mvmhost.Result, err error) { + var evalErr *mvmhost.EvalError + require.ErrorAs(t, err, &evalErr) + assert.Equal(t, 3, evalErr.ExitCode) + }}, + } + for _, test := range cases { + worker = startedProcess(t, "exit 0") + test.check(finishResult(t.Context(), worker, io.NopCloser(strings.NewReader("")), &test.message)) + } +} + +func TestWorkerProcessErrors(t *testing.T) { + t.Parallel() + + worker := startedProcess(t, "exit 4") + require.ErrorContains(t, waitForWorker(t.Context(), worker), "wait for execute worker") + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + worker = startedProcess(t, "exit 4") + require.ErrorIs(t, waitForWorker(ctx, worker), context.Canceled) + + worker = startedProcess(t, "sleep 10") + require.Error(t, worker.abort(errors.New("cause"))) + worker.kill() + + worker = startedProcess(t, "sleep 10") + require.ErrorContains(t, worker.readError(t.Context(), errors.New("read")), "read execute worker") + worker = startedProcess(t, "exit 2") + require.ErrorContains(t, worker.readError(t.Context(), io.EOF), "exited without result") + worker = startedProcess(t, "exit 0") + require.EqualError(t, worker.readError(t.Context(), io.EOF), "execute worker exited without result") + + worker = startedProcess(t, "sleep 10") + assert.ErrorIs(t, worker.readError(ctx, io.EOF), context.Canceled) +} + +func newRPCCaller() *rpcCaller { + return &rpcCaller{ + in: nil, out: nil, terminalErr: nil, pending: make(map[uint64]chan Message), + nextID: 0, mu: sync.Mutex{}, writeMu: sync.Mutex{}, + } +} + +func newClient(handler RPCHandler, executable string) Client { + return Client{Handler: handler, Executable: executable} +} + +func requestWithArguments(arguments any) *Request { + return &Request{ + Arguments: arguments, Mode: "", Name: "", Source: "", + } +} + +func messageWithID(id uint64) *Message { + message := newMessage("") + message.ID = id + + return &message +} + +func messageWithValue(kind string, value json.RawMessage) *Message { + message := newMessage("") + message.ValueKind = kind + message.Value = value + + return &message +} + +func errorMessage(message, kind string, exitCode int) *Message { + result := newMessage("") + result.Error = message + result.ErrorKind = kind + result.ExitCode = exitCode + + return &result +} + +func mvmResult(value any) mvmhost.Result { + return mvmhost.Result{Value: value, ValueKind: "", Stdout: "", Stderr: ""} +} + +func emptyToolCallResult() ToolCallResult { + return ToolCallResult{Details: nil, Error: "", Content: nil, IsError: false} +} diff --git a/internal/executeworker/protocol.go b/internal/executeworker/protocol.go new file mode 100644 index 00000000..a4521947 --- /dev/null +++ b/internal/executeworker/protocol.go @@ -0,0 +1,160 @@ +// Package executeworker implements the bounded framed protocol used to isolate +// provider-facing MVM execution in a helper process. +package executeworker + +import ( + "encoding/binary" + "encoding/json" + "errors" + "fmt" + "io" + "math" + + "github.com/omarluq/librecode/internal/executionlimits" + "github.com/omarluq/librecode/internal/tool" +) + +const ( + // MaxFrameSize is the largest encoded protocol message accepted by a worker. + MaxFrameSize = executionlimits.MaxFrameSize + // MaxResultSize is the largest encoded evaluation value accepted across the protocol. + MaxResultSize = executionlimits.MaxResultSize +) + +// ToolCallResult is the typed result of a tools.Call callback. Keeping this +// envelope typed on both sides of the worker boundary preserves content blocks +// and details instead of decoding them through map[string]any. +type ToolCallResult struct { + Details map[string]any `json:"details"` + Error string `json:"error,omitempty"` + Content []tool.ContentBlock `json:"content"` + IsError bool `json:"is_error"` +} + +// Message is a framed request or response exchanged with an execute worker. +type Message struct { + Stderr string `json:"stderr,omitempty"` + Source string `json:"source,omitempty"` + Method string `json:"method,omitempty"` + Mode string `json:"mode,omitempty"` + Name string `json:"name,omitempty"` + Query string `json:"query,omitempty"` + Stdout string `json:"stdout,omitempty"` + Type string `json:"type"` + Error string `json:"error,omitempty"` + ErrorKind string `json:"error_kind,omitempty"` + ValueKind string `json:"value_kind,omitempty"` + Input json.RawMessage `json:"input,omitempty"` + Value json.RawMessage `json:"value,omitempty"` + Arguments json.RawMessage `json:"arguments,omitempty"` + ID uint64 `json:"id,omitempty"` + ExitCode int `json:"exit_code,omitempty"` +} + +const ( + toolCallResultKind = "tool_call_result" + pipelineResultKind = "pipeline_result" + jsonNullValue = "null" +) + +func newMessage(messageType string) Message { + return Message{ + Stderr: "", Source: "", Method: "", Mode: "", Name: "", Query: "", Stdout: "", + Type: messageType, Error: "", ErrorKind: "", ValueKind: "", Input: nil, Value: nil, + Arguments: nil, ID: 0, ExitCode: 0, + } +} + +// Read decodes one length-prefixed message from r. +func Read(reader io.Reader) (Message, error) { + var size [4]byte + if _, err := io.ReadFull(reader, size[:]); err != nil { + return Message{}, fmt.Errorf("read execute worker frame size: %w", err) + } + + frameSize := binary.BigEndian.Uint32(size[:]) + if frameSize == 0 || frameSize > MaxFrameSize { + return Message{}, fmt.Errorf("execute worker frame size %d exceeds limit %d", frameSize, MaxFrameSize) + } + + payload := make([]byte, frameSize) + if _, err := io.ReadFull(reader, payload); err != nil { + return Message{}, fmt.Errorf("read execute worker frame: %w", err) + } + + var message Message + if err := json.Unmarshal(payload, &message); err != nil { + return Message{}, fmt.Errorf("decode execute worker frame: %w", err) + } + + if err := validateMessage(&message); err != nil { + return Message{}, err + } + + return message, nil +} + +// Write encodes one length-prefixed message to w. +func Write(writer io.Writer, message *Message) error { + if err := validateMessage(message); err != nil { + return err + } + + payload, err := json.Marshal(message) + if err != nil { + return fmt.Errorf("encode execute worker frame: %w", err) + } + + if len(payload) > MaxFrameSize { + return fmt.Errorf("execute worker frame size %d exceeds limit %d", len(payload), MaxFrameSize) + } + + if len(payload) > math.MaxUint32 { + return fmt.Errorf("execute worker frame size %d exceeds uint32", len(payload)) + } + + var encodedSize [8]byte + binary.BigEndian.PutUint64(encodedSize[:], uint64(len(payload))) + + var size [4]byte + copy(size[:], encodedSize[4:]) + + if err := writeFull(writer, size[:]); err != nil { + return fmt.Errorf("write execute worker frame size: %w", err) + } + + if err := writeFull(writer, payload); err != nil { + return fmt.Errorf("write execute worker frame: %w", err) + } + + return nil +} + +func validateMessage(message *Message) error { + if message == nil || message.Type == "" { + return errors.New("invalid execute worker message: missing type") + } + + if message.Type == "result" && len(message.Value) > MaxResultSize { + return fmt.Errorf("execute worker result size %d exceeds limit %d", len(message.Value), MaxResultSize) + } + + return nil +} + +func writeFull(w io.Writer, payload []byte) error { + for len(payload) > 0 { + written, err := w.Write(payload) + if err != nil { + return fmt.Errorf("write bytes: %w", err) + } + + if written <= 0 { + return io.ErrShortWrite + } + + payload = payload[written:] + } + + return nil +} diff --git a/internal/executeworker/protocol_test.go b/internal/executeworker/protocol_test.go new file mode 100644 index 00000000..82cb3d6c --- /dev/null +++ b/internal/executeworker/protocol_test.go @@ -0,0 +1,82 @@ +package executeworker_test + +import ( + "bytes" + "encoding/binary" + "strings" + "testing" + + "github.com/omarluq/librecode/internal/executeworker" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestProtocolRejectsOversizedMessagesOnReadAndWrite(t *testing.T) { + t.Parallel() + + tests := []struct { + write func(*bytes.Buffer) error + name string + want string + }{ + { + name: "result write", + write: func(buffer *bytes.Buffer) error { + return executeworker.Write(buffer, protocolMessage( + "result", []byte(`"`+strings.Repeat("x", executeworker.MaxResultSize)+`"`), + )) + }, + want: "result size", + }, + { + name: "frame write", + write: func(buffer *bytes.Buffer) error { + return executeworker.Write(buffer, protocolMessage( + "rpc_result", []byte(`"`+strings.Repeat("x", executeworker.MaxFrameSize)+`"`), + )) + }, + want: "frame size", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + err := test.write(new(bytes.Buffer)) + require.Error(t, err) + assert.Contains(t, err.Error(), test.want) + }) + } + + var input bytes.Buffer + require.NoError(t, binary.Write(&input, binary.BigEndian, uint32(executeworker.MaxFrameSize+1))) + _, err := executeworker.Read(&input) + require.Error(t, err) + assert.Contains(t, err.Error(), "frame size") +} + +func TestProtocolResponseBudgetsCompose(t *testing.T) { + t.Parallel() + + message := protocolMessage("result", []byte(`"`+strings.Repeat("v", executeworker.MaxResultSize-2)+`"`)) + message.Stdout = strings.Repeat("\\", (executeworker.MaxFrameSize-executeworker.MaxResultSize)/(2*6)) + message.Stderr = message.Stdout + + var buffer bytes.Buffer + require.NoError(t, executeworker.Write(&buffer, message)) + got, err := executeworker.Read(&buffer) + require.NoError(t, err) + assert.Equal(t, message.Type, got.Type) + assert.Equal(t, message.Value, got.Value) + assert.Equal(t, message.Stdout, got.Stdout) + assert.Equal(t, message.Stderr, got.Stderr) +} + +func protocolMessage(messageType string, value []byte) *executeworker.Message { + return &executeworker.Message{ + Stderr: "", Source: "", Method: "", Mode: "", Name: "", Query: "", Stdout: "", Type: messageType, + Error: "", ErrorKind: "", ValueKind: "", Input: nil, Value: value, Arguments: nil, + ID: 0, ExitCode: 0, + } +} diff --git a/internal/executeworker/worker.go b/internal/executeworker/worker.go new file mode 100644 index 00000000..5413eedd --- /dev/null +++ b/internal/executeworker/worker.go @@ -0,0 +1,344 @@ +package executeworker + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "sync" + + "github.com/omarluq/librecode/internal/mvmhost" +) + +const errorKey = "error" + +type rpcCaller struct { + in io.Reader + out io.Writer + terminalErr error + pending map[uint64]chan Message + nextID uint64 + mu sync.Mutex + writeMu sync.Mutex +} + +// Serve runs one evaluation. Tool bindings are synchronous callback RPCs to the +// parent; only JSON values cross the process boundary. +func Serve(input io.Reader, output io.Writer) error { + request, err := Read(input) + if err != nil { + return err + } + + if request.Type != "eval" { + return fmt.Errorf("unexpected execute worker message %q", request.Type) + } + + caller := &rpcCaller{ + in: input, out: output, pending: make(map[uint64]chan Message), nextID: 0, terminalErr: nil, + mu: sync.Mutex{}, writeMu: sync.Mutex{}, + } + go caller.readResponses() + + bindings, err := workerBindings(&request, caller) + if err != nil { + return err + } + + result, evalErr := mvmhost.New().Eval(context.Background(), mvmhost.Request{ + Bindings: bindings, Name: request.Name, Source: request.Source, + }) + + response := resultMessage(result, evalErr) + + return Write(output, &response) +} + +func workerBindings(request *Message, caller *rpcCaller) (mvmhost.Bindings, error) { + if request.Mode == "" || request.Mode == "execute" { + return mvmhost.Bindings{"tools": { + "Search": func(query string) any { return caller.call("search", "", query, nil) }, + "Describe": func(name string) any { return caller.call("describe", name, "", nil) }, + "Call": func(name string, input any) any { return caller.call("call", name, "", input) }, + }}, nil + } + + if request.Mode != "workflow" { + return nil, fmt.Errorf("unknown execute worker mode %q", request.Mode) + } + + var arguments map[string]any + if len(request.Arguments) > 0 && string(request.Arguments) != jsonNullValue { + if err := json.Unmarshal(request.Arguments, &arguments); err != nil { + return nil, fmt.Errorf("decode workflow arguments: %w", err) + } + } + + return mvmhost.Bindings{"librecode/workflow": { + "Arguments": arguments, + "Agent": func(prompt string, options ...map[string]any) (any, error) { + return caller.callResult("workflow_agent", "", "", map[string]any{"prompt": prompt, "options": options}) + }, + "Wait": func(taskID string) (any, error) { + return caller.callResult("workflow_wait", taskID, "", nil) + }, + "List": func() (any, error) { return caller.callResult("workflow_list", "", "", nil) }, + "Cancel": func(taskID string) (any, error) { + return caller.callResult("workflow_cancel", taskID, "", nil) + }, + "Pipeline": func(items []any, callback func(any) (any, error), concurrency int) (any, error) { + results, err := workerPipeline(items, callback, concurrency) + + return pipelineValue(results), err + }, + }}, nil +} + +type pipelineValue []map[string]any + +func workerPipeline(items []any, callback func(any) (any, error), concurrency int) ([]map[string]any, error) { + if concurrency <= 0 { + return nil, errors.New("pipeline concurrency must be positive") + } + + if callback == nil { + return nil, errors.New("pipeline callback is required") + } + + results := make([]map[string]any, len(items)) + for index := range results { + results[index] = map[string]any{ + "index": index, "value": nil, "error": "pipeline stopped before item was scheduled", + } + } + + var ( + state struct { + sync.Mutex + next int + stopped bool + } + workers sync.WaitGroup + ) + workers.Add(min(concurrency, len(items))) + + for range min(concurrency, len(items)) { + go runWorkerPipeline(items, callback, results, &state, &workers) + } + + workers.Wait() + + return results, nil +} + +func runWorkerPipeline( + items []any, + callback func(any) (any, error), + results []map[string]any, + state *struct { + sync.Mutex + next int + stopped bool + }, + workers *sync.WaitGroup, +) { + defer workers.Done() + + for { + state.Lock() + if state.stopped || state.next >= len(items) { + state.Unlock() + + return + } + + index := state.next + state.next++ + state.Unlock() + + value, err := callback(items[index]) + + message := "" + if err != nil { + message = err.Error() + + state.Lock() + state.stopped = true + state.Unlock() + } + + results[index] = map[string]any{"index": index, "value": value, "error": message} + } +} + +func (caller *rpcCaller) call(method, name, query string, input any) any { + value, err := caller.callResult(method, name, query, input) + if err != nil { + return rpcError(err.Error()) + } + + return value +} + +func (caller *rpcCaller) callResult(method, name, query string, input any) (any, error) { + raw, err := json.Marshal(input) + if err != nil { + return nil, fmt.Errorf("encode worker RPC input: %w", err) + } + + response, err := caller.exchange(method, name, query, raw) + if err != nil { + return nil, err + } + + value, err := decodeRPCValue(response) + if err != nil { + return nil, err + } + + return value, nil +} + +func (caller *rpcCaller) exchange(method, name, query string, input json.RawMessage) (*Message, error) { + caller.mu.Lock() + if caller.terminalErr != nil { + err := caller.terminalErr + caller.mu.Unlock() + + return nil, err + } + + caller.nextID++ + requestID := caller.nextID + responseCh := make(chan Message, 1) + caller.pending[requestID] = responseCh + caller.mu.Unlock() + + request := newMessage("rpc") + request.ID, request.Method, request.Name = requestID, method, name + request.Query, request.Input = query, input + + caller.writeMu.Lock() + err := Write(caller.out, &request) + caller.writeMu.Unlock() + + if err != nil { + caller.mu.Lock() + delete(caller.pending, requestID) + caller.mu.Unlock() + + return nil, err + } + + response := <-responseCh + if response.Error != "" { + return nil, errors.New(response.Error) + } + + return &response, nil +} + +func (caller *rpcCaller) readResponses() { + for { + response, err := Read(caller.in) + if err != nil { + caller.failPending(err) + + return + } + + if response.Type != "rpc_result" { + continue + } + + caller.mu.Lock() + responseCh := caller.pending[response.ID] + delete(caller.pending, response.ID) + caller.mu.Unlock() + + if responseCh != nil { + responseCh <- response + } + } +} + +func (caller *rpcCaller) failPending(err error) { + caller.mu.Lock() + if caller.terminalErr == nil { + caller.terminalErr = err + } + + pending := caller.pending + caller.pending = make(map[uint64]chan Message) + caller.mu.Unlock() + + for _, responseCh := range pending { + response := newMessage("rpc_result") + + response.Error = err.Error() + responseCh <- response + } +} + +func decodeRPCValue(response *Message) (any, error) { + if string(response.Value) == jsonNullValue { + return json.RawMessage(jsonNullValue), nil + } + + if response.ValueKind == toolCallResultKind { + var value ToolCallResult + if err := json.Unmarshal(response.Value, &value); err != nil { + return nil, fmt.Errorf("decode tool call result: %w", err) + } + + return value, nil + } + + var value any + if err := json.Unmarshal(response.Value, &value); err != nil { + return nil, fmt.Errorf("decode RPC result: %w", err) + } + + return value, nil +} + +func rpcError(message string) map[string]any { + return map[string]any{errorKey: message, "is_error": true} +} + +func resultMessage(result mvmhost.Result, evalErr error) Message { + response := newMessage("result") + + response.Stdout, response.Stderr = result.Stdout, result.Stderr + if pipeline, ok := result.Value.(pipelineValue); ok { + result.Value = []map[string]any(pipeline) + response.ValueKind = pipelineResultKind + } + + if evalErr != nil { + response.Error = evalErr.Error() + + var normalized *mvmhost.EvalError + if errors.As(evalErr, &normalized) { + response.ErrorKind = string(normalized.Kind) + response.ExitCode = normalized.ExitCode + } + + return response + } + + value, err := json.Marshal(result.Value) + if err != nil { + response.Error = fmt.Sprintf("encode execute result: %v", err) + response.ErrorKind = string(mvmhost.ErrorKindRuntime) + } else { + if _, ok := result.Value.(ToolCallResult); ok { + response.ValueKind = toolCallResultKind + } + + response.Value = value + } + + return response +} diff --git a/internal/executionlimits/limits.go b/internal/executionlimits/limits.go new file mode 100644 index 00000000..3c963520 --- /dev/null +++ b/internal/executionlimits/limits.go @@ -0,0 +1,13 @@ +// Package executionlimits defines fixed safety budgets for isolated MVM execution. +package executionlimits + +const ( + // MaxFrameSize is the largest encoded worker protocol frame. + MaxFrameSize = 8 << 20 + // MaxResultSize is the largest encoded evaluation value in a result frame. + MaxResultSize = 1 << 20 + // responseEnvelopeBudget reserves space for result metadata and JSON framing. + responseEnvelopeBudget = 1 << 20 + // MaxOutputSize leaves worst-case JSON escaping headroom for captured output. + MaxOutputSize = (MaxFrameSize - MaxResultSize - responseEnvelopeBudget) / 6 +) diff --git a/internal/mvmhost/host.go b/internal/mvmhost/host.go new file mode 100644 index 00000000..4935ffaa --- /dev/null +++ b/internal/mvmhost/host.go @@ -0,0 +1,446 @@ +// Package mvmhost provides a narrow adapter around the MVM interpreter. +package mvmhost + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "io/fs" + "os" + "reflect" + "runtime" + + "github.com/mvm-sh/mvm/interp" + "github.com/mvm-sh/mvm/lang/golang" + "github.com/mvm-sh/mvm/vm" + "github.com/samber/oops" + + "github.com/omarluq/librecode/internal/executionlimits" +) + +const ( + defaultSourceLimit = 256 << 10 + defaultOutputLimit = executionlimits.MaxOutputSize +) + +// Bindings maps an import path to the Go values available to interpreted code. +type Bindings map[string]map[string]any + +// Request describes one isolated MVM evaluation. +type Request struct { + Bindings Bindings + Name string + Source string +} + +// Result contains the evaluated value and separately captured diagnostics. +type Result struct { + Value any + ValueKind string + Stdout string + Stderr string +} + +// ErrorKind classifies an evaluation failure. +type ErrorKind string + +// Evaluation error kinds exposed by the host boundary. +const ( + ErrorKindCompile ErrorKind = "compile" + ErrorKindRuntime ErrorKind = "runtime" + ErrorKindPanic ErrorKind = "panic" + ErrorKindExit ErrorKind = "exit" + ErrorKindCanceled ErrorKind = "canceled" + ErrorKindSourceLimit ErrorKind = "source_limit" + ErrorKindOutputLimit ErrorKind = "output_limit" +) + +// EvalError is a normalized MVM evaluation error. +type EvalError struct { + Err error + Kind ErrorKind + ExitCode int +} + +func (e *EvalError) Error() string { + if e.Kind == ErrorKindExit { + return fmt.Sprintf("mvm %s error (exit code %d): %v", e.Kind, e.ExitCode, e.Err) + } + + return fmt.Sprintf("mvm %s error: %v", e.Kind, e.Err) +} + +func (e *EvalError) Unwrap() error { return e.Err } + +// Evaluator executes source using a fresh interpreter for every call. +type Evaluator struct{} + +// New creates an MVM evaluator. +func New() *Evaluator { return &Evaluator{} } + +// Eval evaluates source synchronously. MVM v0.5.0 cannot interrupt arbitrary +// bytecode, so the context is checked before evaluation and by host bindings; +// untrusted source must invoke this adapter in a killable helper process. +func (e *Evaluator) Eval(ctx context.Context, request Request) (Result, error) { + if err := ctx.Err(); err != nil { + return Result{}, evalError(ErrorKindCanceled, 0, err) + } + + if len(request.Source) > defaultSourceLimit { + err := fmt.Errorf("source is %d bytes; limit is %d", len(request.Source), defaultSourceLimit) + + return Result{}, evalError(ErrorKindSourceLimit, 0, err) + } + + outputBudget := &limitedOutputBudget{remaining: defaultOutputLimit, overflow: false} + stdout := newLimitedBuffer(outputBudget) + stderr := newLimitedBuffer(outputBudget) + value, evalErr := evaluate(request, stdout, stderr) + result := Result{Value: nil, ValueKind: "", Stdout: stdout.String(), Stderr: stderr.String()} + + if stdout.Overflowed() || stderr.Overflowed() { + return result, evalError(ErrorKindOutputLimit, 0, errOutputLimit) + } + + if evalErr != nil { + return result, normalizeError(evalErr) + } + + if err := ctx.Err(); err != nil { + return result, evalError(ErrorKindCanceled, 0, err) + } + + return finalizeResult(result, value) +} + +func finalizeResult(result Result, value reflect.Value) (Result, error) { + if !value.IsValid() { + return result, nil + } + + hostValue := value.Interface() + if hostErr, ok := hostValue.(error); ok && hostErr != nil { + return result, evalError(ErrorKindRuntime, 0, hostErr) + } + + if err := validateValue(value, "result"); err != nil { + return result, evalError(ErrorKindRuntime, 0, err) + } + + result.Value = hostValue + + return result, nil +} + +func evaluate(request Request, stdout, stderr io.Writer) (reflect.Value, error) { + machine := interp.NewInterpreter(golang.GoSpec) + machine.SetIO(bytes.NewReader(nil), stdout, stderr) + // os.DevNull is a portable non-directory path, so package lookups cannot + // fall back to the working directory. Disable MVM's embedded source-package + // and remote fallbacks as well; imports are limited to explicit bindings. + machine.SetPkgfs(os.DevNull) + machine.SetStdlibFS(emptyFS{}) + machine.SetRemoteFS(emptyFS{}) + + bindings, err := reflectBindings(request.Bindings) + if err != nil { + return reflect.Value{}, err + } + + machine.ImportPackageValues(bindings) + + name := request.Name + if name == "" { + name = "execute.go" + } + + value, err := machine.Eval(name, request.Source) + if err != nil { + return value, oops.In("mvmhost").Code("mvm_eval").Wrapf(err, "run MVM source") + } + + return value, nil +} + +func reflectBindings(bindings Bindings) (map[string]map[string]reflect.Value, error) { + values := make(map[string]map[string]reflect.Value, len(bindings)) + for packageName, packageBindings := range bindings { + items := make(map[string]reflect.Value, len(packageBindings)) + for name, value := range packageBindings { + reflected := reflect.ValueOf(value) + if !reflected.IsValid() { + return nil, &valueError{message: fmt.Sprintf("binding %q.%q is nil", packageName, name)} + } + + if err := validateBinding(reflected, packageName+"."+name); err != nil { + return nil, err + } + + items[name] = reflected + } + + values[packageName] = items + } + + return values, nil +} + +func validateBinding(value reflect.Value, name string) error { + if value.Kind() == reflect.Func { + if value.IsNil() { + return &valueError{message: fmt.Sprintf("binding %q is a nil function", name)} + } + + return nil + } + + return validateValue(value, "binding "+name) +} + +type valueVisit struct { + typeOfValue reflect.Type + pointer uintptr + length int + capacity int +} + +type valueValidator struct { + active map[valueVisit]string +} + +type valueCategory uint8 + +const ( + valueUnsupported valueCategory = iota + valueScalar + valueInterface + valuePointer + valueArray + valueSlice + valueMap + valueStruct +) + +func validateValue(value reflect.Value, path string) error { + validator := valueValidator{active: make(map[valueVisit]string)} + + return validator.validate(value, path) +} + +func (v *valueValidator) validate(value reflect.Value, path string) error { + if !value.IsValid() { + return &valueError{message: path + " is invalid"} + } + + switch categorizeValue(value.Kind()) { + case valueScalar: + return nil + case valueInterface: + return v.validateIndirect(value, path) + case valuePointer, valueSlice, valueMap: + if value.IsNil() { + return nil + } + + return v.validateReference(value, path) + case valueArray: + return v.validateSequence(value, path) + case valueStruct: + return v.validateStruct(value, path) + case valueUnsupported: + return &valueError{message: fmt.Sprintf("%s has unsupported type %s", path, value.Type())} + default: + return &valueError{message: fmt.Sprintf("%s has unknown type %s", path, value.Type())} + } +} + +func (v *valueValidator) validateReference(value reflect.Value, path string) error { + visit := valueVisit{typeOfValue: value.Type(), pointer: value.Pointer(), length: 0, capacity: 0} + if value.Kind() == reflect.Slice { + visit.length = value.Len() + visit.capacity = value.Cap() + } + + if previousPath, ok := v.active[visit]; ok { + return &valueError{message: fmt.Sprintf("%s contains a cycle through %s", path, previousPath)} + } + + v.active[visit] = path + defer delete(v.active, visit) + + if value.Kind() == reflect.Pointer { + return v.validate(value.Elem(), path) + } + + if value.Kind() == reflect.Slice { + return v.validateSequence(value, path) + } + + if value.Kind() == reflect.Map { + return v.validateMap(value, path) + } + + return &valueError{message: fmt.Sprintf("%s has unsupported reference type %s", path, value.Type())} +} + +func categorizeValue(kind reflect.Kind) valueCategory { + if kind == reflect.Bool || kind == reflect.String || + (kind >= reflect.Int && kind <= reflect.Int64) || + (kind >= reflect.Uint && kind <= reflect.Uint64 && kind != reflect.Uintptr) || + (kind >= reflect.Float32 && kind <= reflect.Float64) { + return valueScalar + } + + categories := map[reflect.Kind]valueCategory{ + reflect.Interface: valueInterface, + reflect.Pointer: valuePointer, + reflect.Array: valueArray, + reflect.Slice: valueSlice, + reflect.Map: valueMap, + reflect.Struct: valueStruct, + } + + return categories[kind] +} + +func (v *valueValidator) validateIndirect(value reflect.Value, path string) error { + if value.IsNil() { + return nil + } + + return v.validate(value.Elem(), path) +} + +func (v *valueValidator) validateSequence(value reflect.Value, path string) error { + for index := range value.Len() { + if err := v.validate(value.Index(index), fmt.Sprintf("%s[%d]", path, index)); err != nil { + return err + } + } + + return nil +} + +func (v *valueValidator) validateMap(value reflect.Value, path string) error { + iterator := value.MapRange() + for iterator.Next() { + if err := v.validate(iterator.Key(), path+" map key"); err != nil { + return err + } + + if err := v.validate(iterator.Value(), path+" map value"); err != nil { + return err + } + } + + return nil +} + +func (v *valueValidator) validateStruct(value reflect.Value, path string) error { + typeOfValue := value.Type() + for index := range value.NumField() { + field := typeOfValue.Field(index) + if !field.IsExported() { + continue + } + + if err := v.validate(value.Field(index), path+"."+field.Name); err != nil { + return err + } + } + + return nil +} + +func normalizeError(err error) error { + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return evalError(ErrorKindCanceled, 0, err) + } + + var exitErr *interp.ExitError + if errors.As(err, &exitErr) { + return evalError(ErrorKindExit, exitErr.Code, err) + } + + var valueErr *valueError + if errors.As(err, &valueErr) { + return evalError(ErrorKindRuntime, 0, err) + } + + var panicErr *vm.PanicError + if errors.As(err, &panicErr) { + if _, ok := panicErr.Raw.(runtime.Error); ok { + return evalError(ErrorKindRuntime, 0, err) + } + + return evalError(ErrorKindPanic, 0, err) + } + // MVM reports parser/compiler failures as ordinary errors. + return evalError(ErrorKindCompile, 0, err) +} + +type valueError struct{ message string } + +func (e *valueError) Error() string { return e.message } + +func evalError(kind ErrorKind, exitCode int, err error) error { + normalized := &EvalError{Err: err, Kind: kind, ExitCode: exitCode} + + return oops.In("mvmhost").Code(string(kind)).Wrapf(normalized, "evaluate MVM source") +} + +type emptyFS struct{} + +func (emptyFS) Open(string) (fs.File, error) { return nil, fs.ErrNotExist } + +var errOutputLimit = errors.New("MVM output limit exceeded") + +type limitedOutputBudget struct { + remaining int + overflow bool +} + +type limitedBuffer struct { + budget *limitedOutputBudget + buffer bytes.Buffer +} + +func newLimitedBuffer(budget *limitedOutputBudget) *limitedBuffer { + return &limitedBuffer{buffer: bytes.Buffer{}, budget: budget} +} + +func (w *limitedBuffer) Write(data []byte) (int, error) { + if len(data) <= w.budget.remaining { + written, err := w.buffer.Write(data) + w.budget.remaining -= written + + if err != nil { + return written, oops.In("mvmhost").Code("output_capture").Wrapf(err, "capture MVM output") + } + + return written, nil + } + + allowed := w.budget.remaining + if allowed > 0 { + written, err := w.buffer.Write(data[:allowed]) + w.budget.remaining -= written + + if err != nil { + return written, oops.In("mvmhost").Code("output_capture").Wrapf(err, "capture bounded MVM output") + } + + allowed = written + } + + w.budget.overflow = true + + return allowed, errOutputLimit +} + +func (w *limitedBuffer) String() string { return w.buffer.String() } +func (w *limitedBuffer) Overflowed() bool { return w.budget.overflow } + +var _ io.Writer = (*limitedBuffer)(nil) diff --git a/internal/mvmhost/host_test.go b/internal/mvmhost/host_test.go new file mode 100644 index 00000000..64616051 --- /dev/null +++ b/internal/mvmhost/host_test.go @@ -0,0 +1,397 @@ +package mvmhost_test + +import ( + "context" + "errors" + "fmt" + "sync" + "testing" + + "github.com/mvm-sh/mvm/interp" + "github.com/omarluq/librecode/internal/mvmhost" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const ( + hostImport = "host" + hostValueBinding = "Value" + hostValuesBinding = "Values" +) + +type testDTO struct { + Name string + Count int +} + +type graphNode struct { + Next *graphNode +} + +type sharedGraph struct { + Left *graphNode + Right *graphNode +} + +type evalTestCase struct { + wantValue any + wantCause error + name string + wantStdout string + wantKind mvmhost.ErrorKind + request mvmhost.Request + wantExitCode int +} + +func newRequest(source string) mvmhost.Request { + return mvmhost.Request{ + Bindings: nil, + Name: "", + Source: source, + } +} + +func requestWithBindings(source string, bindings map[string]any) mvmhost.Request { + request := newRequest(source) + request.Bindings = mvmhost.Bindings{hostImport: bindings} + + return request +} + +func valueCase(name string, request mvmhost.Request, value any) evalTestCase { + return evalTestCase{ + wantValue: value, + wantCause: nil, + name: name, + wantStdout: "", + wantKind: "", + request: request, + wantExitCode: 0, + } +} + +func errorCase(name string, request mvmhost.Request, kind mvmhost.ErrorKind) evalTestCase { + return evalTestCase{ + wantValue: nil, + wantCause: nil, + name: name, + wantStdout: "", + wantKind: kind, + request: request, + wantExitCode: 0, + } +} + +func valueTestCases() []evalTestCase { + stdoutCase := valueCase( + "stdout is separate from value", + newRequest(`println("diagnostic"); 7`), + 7, + ) + stdoutCase.wantStdout = "diagnostic\n" + + return []evalTestCase{ + valueCase("integer expression", newRequest(`1 + 2`), 3), + valueCase("string value", newRequest(`"value"`), "value"), + valueCase("boolean value", newRequest(`true`), true), + valueCase( + "host binding", + requestWithBindings(`import "host"; host.Add(20, 22)`, map[string]any{ + "Add": func(left, right int) int { return left + right }, + }), + 42, + ), + valueCase( + "slice host value", + requestWithBindings(`import "host"; host.Values()`, map[string]any{ + hostValuesBinding: func() []string { return []string{"a", "b"} }, + }), + []string{"a", "b"}, + ), + valueCase( + "map host value", + requestWithBindings(`import "host"; host.Values()`, map[string]any{ + hostValuesBinding: func() map[string]int { return map[string]int{"a": 1} }, + }), + map[string]int{"a": 1}, + ), + valueCase( + "nil host value", + requestWithBindings(`import "host"; host.Value()`, map[string]any{ + hostValueBinding: func() any { return nil }, + }), + nil, + ), + valueCase( + "DTO host value", + requestWithBindings(`import "host"; host.Value()`, map[string]any{ + hostValueBinding: func() testDTO { return testDTO{Name: "item", Count: 2} }, + }), + testDTO{Name: "item", Count: 2}, + ), + stdoutCase, + } +} + +func errorTestCases(hostCause error) []evalTestCase { + hostErrorCase := errorCase( + "host error", + requestWithBindings(`import "host"; host.Fail()`, map[string]any{ + "Fail": func() error { return hostCause }, + }), + mvmhost.ErrorKindRuntime, + ) + hostErrorCase.wantCause = hostCause + + exitCase := errorCase( + "virtualized exit", + requestWithBindings(`import "host"; host.Exit(7)`, map[string]any{ + "Exit": func(code int) { panic(&interp.ExitError{Code: code}) }, + }), + mvmhost.ErrorKindExit, + ) + exitCase.wantExitCode = 7 + + unsupportedBinding := errorCase( + "unsupported binding", + requestWithBindings(`1`, map[string]any{"Values": make(chan int)}), + mvmhost.ErrorKindRuntime, + ) + nilBinding := errorCase( + "nil binding", + requestWithBindings(`1`, map[string]any{hostValueBinding: nil}), + mvmhost.ErrorKindRuntime, + ) + + var nilFunction func() + + nilFunctionBinding := errorCase( + "nil function binding", + requestWithBindings(`1`, map[string]any{"Call": nilFunction}), + mvmhost.ErrorKindRuntime, + ) + + return []evalTestCase{ + errorCase("syntax error", newRequest(`func {`), mvmhost.ErrorKindCompile), + errorCase("runtime error", newRequest(`1 / 0`), mvmhost.ErrorKindRuntime), + errorCase("panic", newRequest(`panic("boom")`), mvmhost.ErrorKindPanic), + hostErrorCase, + exitCase, + unsupportedBinding, + nilBinding, + nilFunctionBinding, + errorCase( + "ambient import denied", + newRequest(`import "example.invalid/ambient"; 1`), + mvmhost.ErrorKindCompile, + ), + errorCase( + "embedded stdlib source import denied", + newRequest(`import "slices"; slices.Clone([]int{1})`), + mvmhost.ErrorKindCompile, + ), + } +} + +func TestEvaluator_Eval(t *testing.T) { + t.Parallel() + + hostCause := errors.New("host failure") + tests := append(valueTestCases(), errorTestCases(hostCause)...) + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + result, err := mvmhost.New().Eval(context.Background(), test.request) + assert.Equal(t, test.wantStdout, result.Stdout) + + if test.wantKind == "" { + require.NoError(t, err) + assert.Equal(t, test.wantValue, result.Value) + + return + } + + require.Error(t, err) + + var evalErr *mvmhost.EvalError + require.ErrorAs(t, err, &evalErr) + assert.Equal(t, test.wantKind, evalErr.Kind) + assert.Equal(t, test.wantExitCode, evalErr.ExitCode) + + if test.wantCause != nil { + require.ErrorIs(t, err, test.wantCause) + } + }) + } +} + +func TestEvaluator_EvalRejectsCyclicValueGraphs(t *testing.T) { + t.Parallel() + + tests := []struct { + request func() mvmhost.Request + name string + }{ + { + name: "pointer binding", + request: func() mvmhost.Request { + node := &graphNode{Next: nil} + node.Next = node + + return requestWithBindings(`1`, map[string]any{hostValueBinding: node}) + }, + }, + { + name: "map binding", + request: func() mvmhost.Request { + value := make(map[string]any) + value["self"] = value + + return requestWithBindings(`1`, map[string]any{hostValueBinding: value}) + }, + }, + { + name: "slice binding", + request: func() mvmhost.Request { + value := make([]any, 1) + value[0] = value + + return requestWithBindings(`1`, map[string]any{hostValueBinding: value}) + }, + }, + { + name: "host result", + request: func() mvmhost.Request { + return requestWithBindings(`import "host"; host.Value()`, map[string]any{ + hostValueBinding: func() any { + value := make(map[string]any) + value["self"] = value + + return value + }, + }) + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + _, err := mvmhost.New().Eval(context.Background(), test.request()) + require.Error(t, err) + assert.Contains(t, err.Error(), "cycle") + + var evalErr *mvmhost.EvalError + require.ErrorAs(t, err, &evalErr) + assert.Equal(t, mvmhost.ErrorKindRuntime, evalErr.Kind) + }) + } +} + +func TestEvaluator_EvalAllowsSharedAcyclicValues(t *testing.T) { + t.Parallel() + + leaf := &graphNode{Next: nil} + value := sharedGraph{Left: leaf, Right: leaf} + request := requestWithBindings(`1`, map[string]any{hostValueBinding: value}) + + result, err := mvmhost.New().Eval(context.Background(), request) + require.NoError(t, err) + assert.Equal(t, 1, result.Value) +} + +func TestEvaluator_EvalCancellation(t *testing.T) { + t.Parallel() + + t.Run("before start", func(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := mvmhost.New().Eval(ctx, newRequest(`1`)) + require.Error(t, err) + require.ErrorIs(t, err, context.Canceled) + }) + + t.Run("cooperative host call", func(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + started := make(chan struct{}) + request := requestWithBindings(`import "host"; host.Wait()`, map[string]any{ + "Wait": func() error { + close(started) + <-ctx.Done() + + return ctx.Err() + }, + }) + done := make(chan error, 1) + + go func() { + _, err := mvmhost.New().Eval(ctx, request) + done <- err + }() + + <-started + cancel() + + err := <-done + require.Error(t, err) + require.ErrorIs(t, err, context.Canceled) + + var evalErr *mvmhost.EvalError + require.ErrorAs(t, err, &evalErr) + assert.Equal(t, mvmhost.ErrorKindCanceled, evalErr.Kind) + }) +} + +func TestEvaluator_ConcurrentInvocationsAreIsolated(t *testing.T) { + t.Parallel() + + const invocationCount = 16 + + var waitGroup sync.WaitGroup + + evaluator := mvmhost.New() + + errorsByInvocation := make(chan error, invocationCount) + for index := range invocationCount { + waitGroup.Go(func() { + source := fmt.Sprintf("%d + 1", index) + + result, err := evaluator.Eval(context.Background(), newRequest(source)) + if err == nil && result.Value != index+1 { + err = fmt.Errorf("value = %v, want %d", result.Value, index+1) + } + + errorsByInvocation <- err + }) + } + + waitGroup.Wait() + close(errorsByInvocation) + + for err := range errorsByInvocation { + require.NoError(t, err) + } +} + +func TestEvalError_ErrorIncludesZeroExitCode(t *testing.T) { + t.Parallel() + + err := &mvmhost.EvalError{Err: errors.New("exit"), Kind: mvmhost.ErrorKindExit, ExitCode: 0} + assert.Contains(t, err.Error(), "exit code 0") +} + +func TestEvalError_Unwrap(t *testing.T) { + t.Parallel() + + cause := errors.New("cause") + err := &mvmhost.EvalError{Err: cause, Kind: mvmhost.ErrorKindRuntime, ExitCode: 0} + assert.ErrorIs(t, err, cause) +} diff --git a/internal/provider/anthropic.go b/internal/provider/anthropic.go index b6a9e8e3..c16f4749 100644 --- a/internal/provider/anthropic.go +++ b/internal/provider/anthropic.go @@ -53,7 +53,7 @@ func (client *HTTPCompletionClient) advanceAnthropicLoop( return false, err } - state.result.Usage = mergeUsage(state.result.Usage, providerResult.Usage) + state.result.Usage = accumulateUsage(state.result.Usage, providerResult.Usage) if validateErr := validateToolCalls(providerResult.ToolCalls); validateErr != nil { return false, validateErr } diff --git a/internal/provider/openai_chat.go b/internal/provider/openai_chat.go index 2878ed4d..b65a3bc8 100644 --- a/internal/provider/openai_chat.go +++ b/internal/provider/openai_chat.go @@ -62,7 +62,7 @@ func (client *HTTPCompletionClient) advanceOpenAIChatLoop( return false, err } - state.result.Usage = mergeUsage(state.result.Usage, providerResult.Usage) + state.result.Usage = accumulateUsage(state.result.Usage, providerResult.Usage) if validateErr := validateToolCalls(providerResult.ToolCalls); validateErr != nil { return false, validateErr } diff --git a/internal/provider/openai_responses.go b/internal/provider/openai_responses.go index 9f6e22cb..a4bc38b0 100644 --- a/internal/provider/openai_responses.go +++ b/internal/provider/openai_responses.go @@ -57,7 +57,7 @@ func (client *HTTPCompletionClient) completeResponsesLoop( } appendThinking(result, providerResult.Thinking) - result.Usage = mergeUsage(result.Usage, providerResult.Usage) + result.Usage = accumulateUsage(result.Usage, providerResult.Usage) result.FinishReason = providerResult.FinishReason if validateErr := validateToolCalls(providerResult.ToolCalls); validateErr != nil { diff --git a/internal/provider/usage.go b/internal/provider/usage.go index 0df6aa80..11dd8bbd 100644 --- a/internal/provider/usage.go +++ b/internal/provider/usage.go @@ -7,24 +7,15 @@ import ( ) func mergeUsage(estimated, reported llm.Usage) llm.Usage { - usage := estimated - if reported.ContextWindow > 0 { - usage.ContextWindow = reported.ContextWindow - } - - if reported.ContextTokens > 0 { - usage.ContextTokens = reported.ContextTokens - } - - if reported.InputTokens > 0 { - usage.InputTokens = reported.InputTokens - } + return llm.MergeUsage(estimated, reported) +} - if reported.OutputTokens > 0 { - usage.OutputTokens = reported.OutputTokens - } +func accumulateUsage(aggregate, reported llm.Usage) llm.Usage { + usage := llm.MergeUsage(aggregate, reported) + usage.InputTokens = aggregate.InputTokens + reported.InputTokens + usage.OutputTokens = aggregate.OutputTokens + reported.OutputTokens - return llm.MergeUsage(usage, reported) + return usage } func usageFromObject(value any) llm.Usage { diff --git a/internal/provider/usage_merge_internal_test.go b/internal/provider/usage_merge_internal_test.go index bdc62135..f2e6050b 100644 --- a/internal/provider/usage_merge_internal_test.go +++ b/internal/provider/usage_merge_internal_test.go @@ -43,6 +43,25 @@ func TestMergeUsageKeepsExistingBreakdownWhenPresent(t *testing.T) { assert.Equal(t, 5, merged.OutputTokens) } +func TestMergeUsageAccumulatesRequestTokens(t *testing.T) { + t.Parallel() + + merged := accumulateUsage( + llm.Usage{ + Breakdown: nil, TopContributors: nil, ContextWindow: 100, ContextTokens: 20, + InputTokens: 7, OutputTokens: 3, + }, + llm.Usage{ + Breakdown: nil, TopContributors: nil, ContextWindow: 100, ContextTokens: 30, + InputTokens: 11, OutputTokens: 5, + }, + ) + + assert.Equal(t, 18, merged.InputTokens) + assert.Equal(t, 8, merged.OutputTokens) + assert.Equal(t, 30, merged.ContextTokens) +} + func TestUsageFromObjectIgnoresNonObjectsAndMissingValues(t *testing.T) { t.Parallel() diff --git a/internal/terminal/agent_tasks.go b/internal/terminal/agent_tasks.go index da928c36..9206450d 100644 --- a/internal/terminal/agent_tasks.go +++ b/internal/terminal/agent_tasks.go @@ -26,6 +26,7 @@ const ( ) const ( + workflowToolName = "workflow" agentStartToolName = "agent_start" agentStatusToolName = "agent_status" agentWaitToolName = "agent_wait" @@ -78,6 +79,10 @@ func (app *App) trackStartedAgentTask(ctx context.Context, event *assistant.Tool return } + if task.Task.ParentTaskID != "" { + return + } + if isTerminalAgentTaskState(task.Task.State) { app.deliverAgentTaskCompletion(ctx, task) @@ -102,6 +107,7 @@ func agentTaskIDFromDetails(detailsJSON string) string { func (app *App) resetAgentTaskTracking() { app.stopAgentTaskWatches() app.agentTasks = nil + app.activeWorkflows = nil app.agentTasksRefreshedAt = time.Time{} app.deliveredAgentTasks = map[string]struct{}{} } @@ -109,6 +115,7 @@ func (app *App) resetAgentTaskTracking() { func (app *App) refreshVisibleAgentTasks(ctx context.Context) { if app.runtime == nil || app.sessionID == "" { app.agentTasks = nil + app.activeWorkflows = nil return } @@ -119,9 +126,165 @@ func (app *App) refreshVisibleAgentTasks(ctx context.Context) { app.refreshActiveAgentTasks(ctx) } + app.refreshActiveWorkflows(ctx) + app.agentTasksRefreshedAt = time.Now() } +func (app *App) refreshActiveWorkflows(ctx context.Context) { + if app.workflows == nil { + app.activeWorkflows = nil + + return + } + + runs, err := app.workflows.List(ctx, app.sessionID, agentTaskInlineLimit) + if err != nil { + return + } + + listed := make(map[string]database.WorkflowRunEntity, len(runs)) + for index := range runs { + listed[runs[index].Task.ID] = runs[index] + } + + active := make([]database.WorkflowRunEntity, 0, len(runs)) + + for index := range app.activeWorkflows { + previous := app.activeWorkflows[index] + + latest, keep := app.reconcileActiveWorkflow(ctx, &previous, listed) + delete(listed, previous.Task.ID) + + if !keep { + continue + } + + if isTerminalAgentTaskState(latest.Task.State) { + app.deliverWorkflowFailure(ctx, &latest) + + continue + } + + active = append(active, latest) + } + + for index := range runs { + run, found := listed[runs[index].Task.ID] + if found && !isTerminalAgentTaskState(run.Task.State) { + active = append(active, run) + } + } + + app.activeWorkflows = active +} + +func (app *App) reconcileActiveWorkflow( + ctx context.Context, + previous *database.WorkflowRunEntity, + listed map[string]database.WorkflowRunEntity, +) (database.WorkflowRunEntity, bool) { + if latest, found := listed[previous.Task.ID]; found { + return latest, true + } + + loaded, found, err := app.workflows.Get(ctx, previous.Task.ID) + if err != nil { + return *previous, true + } + + if !found { + return *previous, false + } + + return *loaded, true +} + +func (app *App) trackStartedWorkflow(ctx context.Context, event *assistant.ToolEvent) { + runID := workflowRunIDFromDetails(event.DetailsJSON) + if runID == "" || app.workflows == nil { + app.refreshActiveWorkflows(ctx) + + return + } + + for index := range app.activeWorkflows { + if app.activeWorkflows[index].Task.ID == runID { + return + } + } + + run, found, err := app.workflows.Get(ctx, runID) + if err != nil || !found || run.Task.OwnerSessionID != app.sessionID { + app.refreshActiveWorkflows(ctx) + + return + } + + if isTerminalAgentTaskState(run.Task.State) { + app.deliverWorkflowFailure(ctx, run) + + return + } + + app.activeWorkflows = append(app.activeWorkflows, *run) +} + +func workflowRunIDFromDetails(detailsJSON string) string { + var details struct { + RunID string `json:"run_id"` + } + if json.Unmarshal([]byte(detailsJSON), &details) != nil { + return "" + } + + return strings.TrimSpace(details.RunID) +} + +func (app *App) deliverWorkflowFailure(ctx context.Context, run *database.WorkflowRunEntity) { + if run == nil || run.Task.State != database.TaskFailed { + return + } + + runID := run.Task.ID + if _, delivered := app.deliveredAgentTasks[runID]; delivered { + return + } + + app.deliveredAgentTasks[runID] = struct{}{} + + detail := strings.TrimSpace(run.Task.ErrorMessage) + if detail == "" { + detail = "No error detail was returned." + } + + name := strings.TrimSpace(run.Name) + if name == "" { + name = toolDisplayWorkflow + } + + completion := fmt.Sprintf("Workflow %q (%s) failed.\n\n%s", name, runID, detail) + + app.setStatus("workflow failed") + + content := formatToolEventForUI(&assistant.ToolEvent{ + CallID: "", ParentCallID: "", Sequence: 0, Name: "workflow_result", + ArgumentsJSON: "", DetailsJSON: "", Result: completion, Error: detail, IsError: true, + }) + app.addAgentCompletionMessage(content) + app.persistAgentCompletion(ctx, content) + + prompt := completion + "\n\nA background workflow failed after it was submitted. " + + "Report the failure and relevant next step to the user." + if app.busy() { + app.queuePrompt(prompt, false) + + return + } + + app.sendPromptHidden(ctx, prompt) +} + func (app *App) discoverActiveAgentTasks(ctx context.Context) { tasks, err := app.runtime.AgentTasks(ctx, app.sessionID, agentTaskInlineLimit) if err != nil { @@ -130,7 +293,7 @@ func (app *App) discoverActiveAgentTasks(ctx context.Context) { active := make([]database.AgentTaskEntity, 0, len(tasks)) for index := range tasks { - if isTerminalAgentTaskState(tasks[index].Task.State) { + if tasks[index].Task.ParentTaskID != "" || isTerminalAgentTaskState(tasks[index].Task.State) { continue } @@ -225,18 +388,18 @@ func (app *App) refreshActiveAgentTasks(ctx context.Context) { return } - activeByID := make(map[string]database.AgentTaskEntity, len(tasks)) - for index := range tasks { - if !isTerminalAgentTaskState(tasks[index].Task.State) { - activeByID[tasks[index].Task.ID] = tasks[index] - } - } + activeByID := activeIndependentAgentTasksByID(tasks) active := make([]database.AgentTaskEntity, 0, len(activeByID)) completed := make([]database.AgentTaskEntity, 0) for index := range app.agentTasks { previous := app.agentTasks[index] + if previous.Task.ParentTaskID != "" { + app.stopAgentTaskWatch(previous.Task.ID) + + continue + } task, found := activeByID[previous.Task.ID] if !found { @@ -269,6 +432,20 @@ func (app *App) refreshActiveAgentTasks(ctx context.Context) { } } +func activeIndependentAgentTasksByID( + tasks []database.AgentTaskEntity, +) map[string]database.AgentTaskEntity { + activeByID := make(map[string]database.AgentTaskEntity, len(tasks)) + for index := range tasks { + task := tasks[index] + if task.Task.ParentTaskID == "" && !isTerminalAgentTaskState(task.Task.State) { + activeByID[task.Task.ID] = task + } + } + + return activeByID +} + func (app *App) reconcileMissingAgentTask( ctx context.Context, previous *database.AgentTaskEntity, @@ -324,6 +501,12 @@ func (app *App) deliverAgentTaskCompletion(ctx context.Context, task *database.A return } + if task.Task.ParentTaskID != "" { + app.discardAgentTaskCompletion(task.Task.ID) + + return + } + completion, completed := agentTaskCompletion(database.TaskQueued, task) if !completed { return @@ -341,6 +524,37 @@ func (app *App) deliverAgentTaskCompletionText(ctx context.Context, taskID, comp return } + workflowChild := app.isTrackedWorkflowChild(taskID) || app.isPersistedWorkflowChild(ctx, taskID) + app.discardAgentTaskCompletion(taskID) + + if workflowChild { + return + } + + app.deliverAgentTaskCompletions(ctx, []string{completion}) +} + +func (app *App) isTrackedWorkflowChild(taskID string) bool { + for index := range app.agentTasks { + if app.agentTasks[index].Task.ID == taskID { + return app.agentTasks[index].Task.ParentTaskID != "" + } + } + + return false +} + +func (app *App) isPersistedWorkflowChild(ctx context.Context, taskID string) bool { + if app.runtime == nil { + return false + } + + task, found, err := app.runtime.AgentTask(ctx, taskID) + + return err == nil && found && task.Task.ParentTaskID != "" +} + +func (app *App) discardAgentTaskCompletion(taskID string) { app.deliveredAgentTasks[taskID] = struct{}{} app.stopAgentTaskWatch(taskID) @@ -352,8 +566,6 @@ func (app *App) deliverAgentTaskCompletionText(ctx context.Context, taskID, comp } app.agentTasks = active - - app.deliverAgentTaskCompletions(ctx, []string{completion}) } func (app *App) deliverAgentTaskCompletions(ctx context.Context, completions []string) { @@ -427,6 +639,9 @@ func (app *App) persistAgentCompletion(ctx context.Context, content string) { func formatAgentCompletionForUI(completion string) string { return formatToolEventForUI(&assistant.ToolEvent{ + CallID: "", + ParentCallID: "", + Sequence: 0, Name: "agent_result", ArgumentsJSON: "", DetailsJSON: "", @@ -437,8 +652,15 @@ func formatAgentCompletionForUI(completion string) string { } func (app *App) hasRunningAgentTasks() bool { + for index := range app.activeWorkflows { + if !isTerminalAgentTaskState(app.activeWorkflows[index].Task.State) { + return true + } + } + for index := range app.agentTasks { - if !isTerminalAgentTaskState(app.agentTasks[index].Task.State) { + if app.agentTasks[index].Task.ParentTaskID == "" && + !isTerminalAgentTaskState(app.agentTasks[index].Task.State) { return true } } @@ -458,16 +680,32 @@ func isTerminalAgentTaskState(state database.TaskState) bool { } func (app *App) renderAgentTaskSummary(width int) []tui.Line { - if len(app.agentTasks) == 0 { + if len(app.agentTasks) == 0 && len(app.activeWorkflows) == 0 { return nil } indicatorStyle := tcell.StyleDefault.Foreground(defaultWorkingShimmerBrightColor()).Bold(true) labelStyle := tcell.StyleDefault.Foreground(app.theme.colors[colorMuted]) - lines := make([]tui.Line, 0, len(app.agentTasks)+1) + lines := make([]tui.Line, 0, len(app.activeWorkflows)+len(app.agentTasks)+1) + for index := range app.activeWorkflows { + label := workflowSummaryLabel(&app.activeWorkflows[index]) + line := tui.Line{ + Text: pendingToolIndicator + " " + label, Style: labelStyle, + Spans: []tui.Span{ + {Text: pendingToolIndicator, Style: indicatorStyle}, + {Text: " " + label, Style: labelStyle}, + }, + } + lines = append(lines, line.Truncate(max(1, width))) + } + for index := range app.agentTasks { task := &app.agentTasks[index] + if task.Task.ParentTaskID != "" { + continue + } + label := agentTaskSummaryLabel(task) line := tui.Line{ Text: pendingToolIndicator + " " + label, @@ -480,11 +718,28 @@ func (app *App) renderAgentTaskSummary(width int) []tui.Line { lines = append(lines, line.Truncate(max(1, width))) } + if len(lines) == 0 { + return nil + } + lines = append(lines, tui.NewLine(tcell.StyleDefault, "")) return lines } +func workflowSummaryLabel(run *database.WorkflowRunEntity) string { + if run == nil { + return "workflow" + } + + name := strings.Join(strings.Fields(run.Name), " ") + if name == "" { + return "workflow" + } + + return "workflow(" + name + ")" +} + func agentTaskSummaryLabel(task *database.AgentTaskEntity) string { if task == nil { return agentDefaultDisplayName diff --git a/internal/terminal/agent_tasks_behavior_internal_test.go b/internal/terminal/agent_tasks_behavior_internal_test.go index e9502a56..fad08eb6 100644 --- a/internal/terminal/agent_tasks_behavior_internal_test.go +++ b/internal/terminal/agent_tasks_behavior_internal_test.go @@ -65,6 +65,9 @@ func newAgentTaskControllerStub( func agentToolEvent(name, details string, isError bool) *assistant.ToolEvent { return &assistant.ToolEvent{ + CallID: "", + ParentCallID: "", + Sequence: 0, Name: name, ArgumentsJSON: "", DetailsJSON: details, @@ -262,6 +265,43 @@ func TestAgentTaskCompletionFallbacksAndDelivery(t *testing.T) { app.deliverAgentTaskCompletion(t.Context(), nil) } +func TestWorkflowChildAgentTasksAreHiddenAndNotDelivered(t *testing.T) { + t.Parallel() + + child := behaviorAgentTask("workflow-child", database.TaskRunning) + child.Task.ParentTaskID = "workflow-run" + stub := newAgentTaskControllerStub( + map[string]*database.AgentTaskEntity{child.Task.ID: &child}, + []database.AgentTaskEntity{child}, + ) + app := newAgentTaskBehaviorApp(t, stub) + app.working = true + + app.discoverActiveAgentTasks(t.Context()) + assert.Empty(t, app.agentTasks) + app.activeWorkflows = []database.WorkflowRunEntity{ + workflowSummaryRun("workflow-run", database.TaskRunning), + } + lines := app.renderAgentTaskSummary(80) + require.Len(t, lines, 2) + assert.Equal(t, "â—Œ workflow(active review)", lines[0].Text) + assert.True(t, app.hasRunningAgentTasks()) + + app.trackStartedAgentTask(t.Context(), agentToolEvent("", `{"task_id":"workflow-child"}`, false)) + assert.Empty(t, app.agentTasks) + + completed := child + completed.Task.State = database.TaskSucceeded + completed.Task.Result = "workflow-owned result" + stub.tasks[child.Task.ID] = &completed + app.deliverAgentTaskCompletionText(t.Context(), child.Task.ID, "workflow-owned result") + app.deliverAgentTaskCompletion(t.Context(), &completed) + + assert.Empty(t, app.liveAgentCompletions) + assert.Empty(t, app.hiddenQueuedMessages) + assert.Contains(t, app.deliveredAgentTasks, child.Task.ID) +} + func TestDiscoverRefreshAndTrackAgentTasks(t *testing.T) { t.Parallel() @@ -605,7 +645,7 @@ func TestAgentTaskPanelRejectsTaskOwnedByAnotherSession(t *testing.T) { t.Parallel() task := behaviorAgentTask(behaviorTaskID, database.TaskRunning) - task.Task.OwnerSessionID = "another-session" + task.Task.OwnerSessionID = workflowTestForeignSession stub := newAgentTaskControllerStub(map[string]*database.AgentTaskEntity{behaviorTaskID: &task}, nil) app := newAgentTaskBehaviorApp(t, stub) items := []tui.ListItem{{Value: behaviorTaskID, Title: "task", Description: "", Meta: ""}} diff --git a/internal/terminal/app.go b/internal/terminal/app.go index 1281e66d..fce0e95e 100644 --- a/internal/terminal/app.go +++ b/internal/terminal/app.go @@ -22,6 +22,11 @@ import ( "github.com/omarluq/librecode/internal/tui" ) +type workflowInspector interface { + Get(context.Context, string) (*database.WorkflowRunEntity, bool, error) + List(context.Context, string, int) ([]database.WorkflowRunEntity, error) +} + const ( defaultEditorRows = 6 workFrameInterval = 120 * time.Millisecond @@ -101,6 +106,7 @@ type RunOptions struct { Extensions extension.TerminalEventRunner `json:"-"` Resources *core.ResourceSnapshot `json:"resources"` Runtime *assistant.Runtime `json:"-"` + Workflows workflowInspector `json:"-"` Settings *database.DocumentRepository `json:"-"` Models *model.Registry `json:"-"` Auth *auth.Storage `json:"-"` @@ -121,6 +127,7 @@ type App struct { lastResize *tcell.EventResize systemClipboard systemClipboardWriter runtime *assistant.Runtime + workflows workflowInspector settings *database.DocumentRepository models *model.Registry auth *auth.Storage @@ -147,6 +154,7 @@ type App struct { runningToolBlocks []runningToolBlock liveAgentCompletions []chatMessage agentTasks []database.AgentTaskEntity + activeWorkflows []database.WorkflowRunEntity agentTasksRefreshedAt time.Time agentTaskWatches map[string]context.CancelFunc deliveredAgentTasks map[string]struct{} @@ -232,6 +240,7 @@ func newApp(screen terminalScreen, options *RunOptions) *App { lastResize: nil, systemClipboard: newDesktopClipboard(), runtime: options.Runtime, + workflows: options.Workflows, extensions: options.Extensions, settings: options.Settings, models: options.Models, @@ -261,6 +270,7 @@ func newApp(screen terminalScreen, options *RunOptions) *App { runningToolBlocks: []runningToolBlock{}, liveAgentCompletions: []chatMessage{}, agentTasks: []database.AgentTaskEntity{}, + activeWorkflows: []database.WorkflowRunEntity{}, agentTasksRefreshedAt: time.Time{}, agentTaskWatches: map[string]context.CancelFunc{}, deliveredAgentTasks: map[string]struct{}{}, diff --git a/internal/terminal/async_events.go b/internal/terminal/async_events.go index 251dba2f..b7df0149 100644 --- a/internal/terminal/async_events.go +++ b/internal/terminal/async_events.go @@ -615,8 +615,10 @@ func (app *App) applyStreamedToolStart(call *assistant.ToolCallEvent, fallbackNa call = &assistant.ToolCallEvent{ Arguments: tool.EmptyArguments(), ID: "", + ParentCallID: "", Name: fallbackName, ArgumentsJSON: "", + Sequence: 0, } } @@ -636,6 +638,10 @@ func (app *App) applyStreamedToolEvent(event *assistant.ToolEvent) { return } + if event.Name == workflowToolName && !event.IsError { + app.trackStartedWorkflow(context.Background(), event) + } + if isAgentManagementTool(event.Name) { app.applyAgentToolEvent(event) diff --git a/internal/terminal/async_events_internal_test.go b/internal/terminal/async_events_internal_test.go index 48cb3618..14a14e26 100644 --- a/internal/terminal/async_events_internal_test.go +++ b/internal/terminal/async_events_internal_test.go @@ -741,6 +741,9 @@ func TestAsyncEventDataHelpers(t *testing.T) { assert.Empty(t, toolExtensionData(nil)) toolData := toolExtensionData(&assistant.ToolEvent{ + CallID: "", + ParentCallID: "", + Sequence: 0, Name: asyncTestToolStart, ArgumentsJSON: "{}", DetailsJSON: "details", @@ -765,6 +768,9 @@ func asyncTestUsage() *model.TokenUsage { func asyncTestToolEvent() *assistant.ToolEvent { return &assistant.ToolEvent{ + CallID: "", + ParentCallID: "", + Sequence: 0, Name: asyncTestToolName, ArgumentsJSON: "{}", DetailsJSON: "", diff --git a/internal/terminal/autocomplete.go b/internal/terminal/autocomplete.go index c013a9b4..04211703 100644 --- a/internal/terminal/autocomplete.go +++ b/internal/terminal/autocomplete.go @@ -374,6 +374,9 @@ func formatToolEventForUI(event *assistant.ToolEvent) string { } return transcript.FormatToolEventDisplay(&transcript.ToolEvent{ + CallID: event.CallID, + ParentCallID: event.ParentCallID, + Sequence: event.Sequence, Name: event.Name, ArgumentsJSON: event.ArgumentsJSON, DetailsJSON: event.DetailsJSON, diff --git a/internal/terminal/autocomplete_internal_test.go b/internal/terminal/autocomplete_internal_test.go index 05dc1040..f2829105 100644 --- a/internal/terminal/autocomplete_internal_test.go +++ b/internal/terminal/autocomplete_internal_test.go @@ -20,6 +20,7 @@ func TestSlashSuggestionsOnlyIncludesImplementedCommands(t *testing.T) { } assert.Contains(t, names, "skill") + assert.NotContains(t, names, "workflows") assert.NotContains(t, names, "export") assert.NotContains(t, names, "import") assert.NotContains(t, names, "share") diff --git a/internal/terminal/bench_terminal_internal_test.go b/internal/terminal/bench_terminal_internal_test.go index e1979877..45672e7a 100644 --- a/internal/terminal/bench_terminal_internal_test.go +++ b/internal/terminal/bench_terminal_internal_test.go @@ -16,6 +16,7 @@ func BenchmarkDrawMessagesSameWidth(b *testing.B) { Extensions: nil, Resources: nil, Runtime: nil, + Workflows: nil, Settings: nil, Models: nil, Auth: nil, diff --git a/internal/terminal/commands_internal_test.go b/internal/terminal/commands_internal_test.go index c5176462..f73bd01c 100644 --- a/internal/terminal/commands_internal_test.go +++ b/internal/terminal/commands_internal_test.go @@ -38,7 +38,7 @@ func TestRenameSessionValidation(t *testing.T) { app := newRenderTestApp(t) require.EqualError(t, app.renameSession(context.Background(), "name"), "no active session") - app.sessionID = "session-1" + app.sessionID = workflowTestSessionID require.EqualError(t, app.renameSession(context.Background(), ""), "name is required") } diff --git a/internal/terminal/message_layout.go b/internal/terminal/message_layout.go index 0e510ead..ec1d69fa 100644 --- a/internal/terminal/message_layout.go +++ b/internal/terminal/message_layout.go @@ -194,7 +194,7 @@ func (app *App) dynamicMessageLineGroups(width int) [][]tui.Line { } for index := range app.runningToolBlocks { - groups = append(groups, app.renderRunningToolBlock(width, app.runningToolBlocks[index].Call)) + groups = append(groups, app.renderRunningToolBlock(width, &app.runningToolBlocks[index].Call)) } if app.busy() { diff --git a/internal/terminal/prompt_cancel_internal_test.go b/internal/terminal/prompt_cancel_internal_test.go index 5c78b81e..2678e45c 100644 --- a/internal/terminal/prompt_cancel_internal_test.go +++ b/internal/terminal/prompt_cancel_internal_test.go @@ -69,8 +69,11 @@ func TestCancelActivePromptPreservesPersistedProgress(t *testing.T) { require.Equal(t, asyncEventPromptUserEntry, userEntryEvent.Kind) app.handlePromptAsyncEvent(context.Background(), userEntryEvent) - request := client.waitForRequest(t) - client.waitForPromptEntry(t) + waitCtx, cancelWait := context.WithTimeout(t.Context(), 5*time.Second) + defer cancelWait() + + request := client.waitForRequest(waitCtx, t) + client.waitForPromptEntry(waitCtx, t) request.OnEvent(assistant.StreamEvent{ ToolCallEvent: nil, ToolEvent: nil, @@ -93,16 +96,18 @@ func TestCancelActivePromptPreservesPersistedProgress(t *testing.T) { } type cancelPreserveCompleter struct { - request *assistant.CompletionRequest - ready chan struct{} - once sync.Once + request *assistant.CompletionRequest + ready chan struct{} + promptReady chan struct{} + once sync.Once } func newCancelPreserveCompleter() *cancelPreserveCompleter { return &cancelPreserveCompleter{ - request: nil, - ready: make(chan struct{}), - once: sync.Once{}, + request: nil, + ready: make(chan struct{}), + promptReady: make(chan struct{}), + once: sync.Once{}, } } @@ -111,39 +116,44 @@ func (completer *cancelPreserveCompleter) Complete( request *assistant.CompletionRequest, ) (*assistant.CompletionResult, error) { completer.request = request - completer.once.Do(func() { close(completer.ready) }) + completer.once.Do(func() { + close(completer.ready) + + if len(request.Messages) > 0 && request.Messages[len(request.Messages)-1].Content == "keep progress" { + close(completer.promptReady) + } + }) <-ctx.Done() return nil, fmt.Errorf("completion canceled: %w", ctx.Err()) } -func (completer *cancelPreserveCompleter) waitForRequest(t *testing.T) *assistant.CompletionRequest { +func (completer *cancelPreserveCompleter) waitForRequest( + ctx context.Context, + t *testing.T, +) *assistant.CompletionRequest { t.Helper() select { case <-completer.ready: + require.NotNil(t, completer.request) + return completer.request - case <-time.After(time.Second): - t.Fatal("timed out waiting for completion request") + case <-ctx.Done(): + require.FailNow(t, "timed out waiting for completion request", ctx.Err().Error()) return nil } } -func (completer *cancelPreserveCompleter) waitForPromptEntry(t *testing.T) { +func (completer *cancelPreserveCompleter) waitForPromptEntry(ctx context.Context, t *testing.T) { t.Helper() - for { - request := completer.request - if request != nil && request.Messages[len(request.Messages)-1].Content == "keep progress" { - return - } - - select { - case <-time.After(10 * time.Millisecond): - case <-time.After(time.Second): - t.Fatal("timed out waiting for prompt entry in completion request") - } + select { + case <-completer.promptReady: + return + case <-ctx.Done(): + require.FailNow(t, "timed out waiting for prompt entry in completion request", ctx.Err().Error()) } } diff --git a/internal/terminal/prompt_response_internal_test.go b/internal/terminal/prompt_response_internal_test.go index 1679b8ee..19134672 100644 --- a/internal/terminal/prompt_response_internal_test.go +++ b/internal/terminal/prompt_response_internal_test.go @@ -68,8 +68,14 @@ func TestApplyRemainingSideEffectsSkipsStreamedBlocks(t *testing.T) { Text: "", Thinking: []string{"streamed thinking", "remaining thinking"}, ToolEvents: []assistant.ToolEvent{ - {Name: "read", ArgumentsJSON: "{}", DetailsJSON: "", Result: "streamed", Error: "", IsError: false}, - {Name: "write", ArgumentsJSON: "{}", DetailsJSON: "", Result: "remaining", Error: "", IsError: false}, + { + CallID: "", ParentCallID: "", Sequence: 0, Name: "read", ArgumentsJSON: "{}", + DetailsJSON: "", Result: "streamed", Error: "", IsError: false, + }, + { + CallID: "", ParentCallID: "", Sequence: 0, Name: "write", ArgumentsJSON: "{}", + DetailsJSON: "", Result: "remaining", Error: "", IsError: false, + }, }, Usage: model.EmptyTokenUsage(), Cached: false, diff --git a/internal/terminal/render_internal_test.go b/internal/terminal/render_internal_test.go index 43fc8e79..327d2b08 100644 --- a/internal/terminal/render_internal_test.go +++ b/internal/terminal/render_internal_test.go @@ -1285,6 +1285,9 @@ func TestBottomMessageLinesExpandedToolUsesTailWithoutFullRender(t *testing.T) { "line 5", }, "\n") app.addMessage(transcript.RoleToolResult, formatToolEventForUI(&assistant.ToolEvent{ + CallID: "", + ParentCallID: "", + Sequence: 0, Name: "read", ArgumentsJSON: "{}", DetailsJSON: "", @@ -1757,6 +1760,7 @@ func newRenderTestApp(t *testing.T) *App { Extensions: nil, Resources: nil, Runtime: nil, + Workflows: nil, Settings: nil, Models: nil, Auth: nil, @@ -1801,6 +1805,9 @@ func newTestAsyncEvent(kind asyncEventKind, text string) *asyncEvent { func newTestToolEvent(name, result string) *assistant.ToolEvent { return &assistant.ToolEvent{ + CallID: "", + ParentCallID: "", + Sequence: 0, Name: name, ArgumentsJSON: "", DetailsJSON: "", diff --git a/internal/terminal/render_parity_internal_test.go b/internal/terminal/render_parity_internal_test.go index 99729476..052b7910 100644 --- a/internal/terminal/render_parity_internal_test.go +++ b/internal/terminal/render_parity_internal_test.go @@ -188,6 +188,7 @@ func TestRenderParityResumedHistoryViewport(t *testing.T) { Extensions: nil, Resources: nil, Runtime: runtime, + Workflows: nil, Settings: nil, Models: nil, Auth: nil, diff --git a/internal/terminal/running_tools.go b/internal/terminal/running_tools.go index 1686d627..5306b859 100644 --- a/internal/terminal/running_tools.go +++ b/internal/terminal/running_tools.go @@ -7,6 +7,14 @@ func (app *App) removeRunningToolBlock(event *assistant.ToolEvent) { return } + if event.CallID != "" { + if index, ok := app.runningToolBlockIndexByCallID(event.CallID); ok { + app.deleteRunningToolBlock(index) + + return + } + } + if index, ok := app.runningToolBlockIndexByArguments(event.Name, event.ArgumentsJSON); ok { app.deleteRunningToolBlock(index) @@ -18,6 +26,20 @@ func (app *App) removeRunningToolBlock(event *assistant.ToolEvent) { } } +func (app *App) runningToolBlockIndexByCallID(callID string) (int, bool) { + if callID == "" { + return 0, false + } + + for index, block := range app.runningToolBlocks { + if block.Call.ID == callID { + return index, true + } + } + + return 0, false +} + func (app *App) runningToolBlockIndexByArguments(name, argumentsJSON string) (int, bool) { if argumentsJSON == "" { return 0, false diff --git a/internal/terminal/running_tools_internal_test.go b/internal/terminal/running_tools_internal_test.go index 65ed7e03..7d2f7a32 100644 --- a/internal/terminal/running_tools_internal_test.go +++ b/internal/terminal/running_tools_internal_test.go @@ -31,6 +31,7 @@ func TestAgentManagementToolsUseTaskSummaryInsteadOfToolBlocks(t *testing.T) { app.applyStreamedToolStart(&call, "") app.applyStreamedToolEvent(&assistant.ToolEvent{ + CallID: "", ParentCallID: "", Sequence: 0, Name: agentStartToolName, ArgumentsJSON: call.ArgumentsJSON, DetailsJSON: "", Result: "started", Error: "", IsError: false, }) @@ -217,7 +218,7 @@ func TestAgentCompletionSurvivesPromptStreamingReset(t *testing.T) { app.appendStreamingBlock(transcript.RoleAssistant, "parent response") app.applyPromptResponse(context.Background(), &assistant.PromptResponse{ - Usage: model.EmptyTokenUsage(), SessionID: "session-1", UserEntryID: "", AssistantEntryID: "", + Usage: model.EmptyTokenUsage(), SessionID: workflowTestSessionID, UserEntryID: "", AssistantEntryID: "", Text: "parent response", Thinking: nil, ToolEvents: nil, Cached: false, }, 1) @@ -294,10 +295,13 @@ func runningToolAppendRenderRemoveCase() runningToolBlockTestCase { app.applyStreamedToolStart(&call, "") require.NotEmpty(t, app.runningToolBlocks) - lines := app.renderRunningToolBlock(80, app.runningToolBlocks[0].Call) + lines := app.renderRunningToolBlock(80, &app.runningToolBlocks[0].Call) assert.NotEqual(t, -1, lineIndexContaining(lines, "â—Œ $ go test ./...")) app.applyStreamedToolEvent(&assistant.ToolEvent{ + CallID: "", + ParentCallID: "", + Sequence: 0, Name: testToolBash, ArgumentsJSON: bashTestArguments, DetailsJSON: "", @@ -333,6 +337,44 @@ func TestApplyStreamedToolStartSkipsAgentFallbackName(t *testing.T) { assert.Empty(t, app.runningToolBlocks) } +func TestRemoveRunningToolBlockPrefersCallID(t *testing.T) { + t.Parallel() + + app := newRenderTestApp(t) + first := testRunningToolBlock(testToolRead, `{"path":"same.go"}`) + first.Call.ID = "call-alpha" + second := testRunningToolBlock(testToolRead, `{"path":"same.go"}`) + second.Call.ID = "call-beta" + app.runningToolBlocks = []runningToolBlock{first, second} + + app.removeRunningToolBlock(&assistant.ToolEvent{ + CallID: "call-beta", ParentCallID: "outer", Name: testToolRead, ArgumentsJSON: `{"path":"same.go"}`, + DetailsJSON: "", Result: "ok", Error: "", Sequence: 2, IsError: false, + }) + + require.Len(t, app.runningToolBlocks, 1) + assert.Equal(t, "call-alpha", app.runningToolBlocks[0].Call.ID) +} + +func TestRemoveRunningToolBlockFallsBackWhenResultHasGeneratedCallID(t *testing.T) { + t.Parallel() + + app := newRenderTestApp(t) + app.runningToolBlocks = []runningToolBlock{ + testRunningToolBlock(testToolRead, `{"path":"a.go"}`), + testRunningToolBlock(testToolWrite, `{"path":"b.go"}`), + } + + app.removeRunningToolBlock(&assistant.ToolEvent{ + Result: "", Error: "", CallID: "generated-call-id", ParentCallID: "", + Name: testToolWrite, ArgumentsJSON: `{"path":"b.go"}`, DetailsJSON: "", + Sequence: 0, IsError: false, + }) + + require.Len(t, app.runningToolBlocks, 1) + assert.Equal(t, testToolRead, app.runningToolBlocks[0].Call.Name) +} + func runningToolRemoveByNameAndArgumentsCase() runningToolBlockTestCase { sharedArguments := `{"path":"same.go"}` @@ -345,6 +387,9 @@ func runningToolRemoveByNameAndArgumentsCase() runningToolBlockTestCase { testRunningToolBlock(testToolWrite, sharedArguments), } app.removeRunningToolBlock(&assistant.ToolEvent{ + CallID: "", + ParentCallID: "", + Sequence: 0, Name: testToolWrite, ArgumentsJSON: sharedArguments, DetailsJSON: "", @@ -368,6 +413,9 @@ func runningToolRemoveByNameFallbackCase() runningToolBlockTestCase { testRunningToolBlock(testToolBash, `{"command":"go test"}`), } app.removeRunningToolBlock(&assistant.ToolEvent{ + CallID: "", + ParentCallID: "", + Sequence: 0, Name: testToolBash, ArgumentsJSON: "", DetailsJSON: "", @@ -414,8 +462,10 @@ func testToolCallEvent(name, argumentsJSON string) assistant.ToolCallEvent { return assistant.ToolCallEvent{ Arguments: tool.EmptyArguments(), ID: "", + ParentCallID: "", Name: name, ArgumentsJSON: argumentsJSON, + Sequence: 0, } } diff --git a/internal/terminal/token_usage_behavior_internal_test.go b/internal/terminal/token_usage_behavior_internal_test.go index 5d50996b..2a24156e 100644 --- a/internal/terminal/token_usage_behavior_internal_test.go +++ b/internal/terminal/token_usage_behavior_internal_test.go @@ -305,6 +305,7 @@ func newTestApp() *App { Extensions: nil, Resources: nil, Runtime: nil, + Workflows: nil, Settings: nil, Models: nil, Auth: nil, diff --git a/internal/terminal/token_usage_export_test.go b/internal/terminal/token_usage_export_test.go index 12885ee4..c06b248c 100644 --- a/internal/terminal/token_usage_export_test.go +++ b/internal/terminal/token_usage_export_test.go @@ -33,6 +33,7 @@ func NewAppForTest() *App { Extensions: nil, Resources: nil, Runtime: nil, + Workflows: nil, Settings: nil, Models: nil, Auth: nil, diff --git a/internal/terminal/tool_blocks.go b/internal/terminal/tool_blocks.go index e719814c..2e4a6591 100644 --- a/internal/terminal/tool_blocks.go +++ b/internal/terminal/tool_blocks.go @@ -16,6 +16,7 @@ const ( maxCollapsedToolOutputLines = 5 toolArgumentDiffContext = 3 toolSectionTool = "tool" + toolSectionParentCallID = "parent_call_id" toolSectionArguments = "arguments" toolSectionDetails = "details" toolSectionError = "error" @@ -26,6 +27,7 @@ const ( type parsedToolEvent struct { Name string + ParentCallID string ArgumentsJSON string DetailsJSON string Error string @@ -39,7 +41,7 @@ func (app *App) renderToolBlock(width int, message chatMessage) []tui.Line { return app.renderToolDisplayBlock(width, &display) } -func (app *App) renderRunningToolBlock(width int, call assistant.ToolCallEvent) []tui.Line { +func (app *App) renderRunningToolBlock(width int, call *assistant.ToolCallEvent) []tui.Line { display := toolDisplayFromCall(call) return app.renderToolDisplayBlock(width, &display) @@ -69,6 +71,14 @@ func (app *App) toolDisplayStyle(display *toolDisplay) tcell.Style { func (app *App) renderPendingToolDisplay(width int, display *toolDisplay, style tcell.Style) []tui.Line { lines := toolBlockStart(width, style) lines = append(lines, app.toolHeaderDisplayLines(width, display, style)...) + + if app.toolsExpanded { + if source := executeSource(display); source != "" { + lines = append(lines, app.toolExpandHintLines(width, style)...) + lines = append(lines, app.toolCodeLines(width, source, style)...) + } + } + lines = append(lines, toolBlockEnd(width, style)...) return lines @@ -107,6 +117,10 @@ func (app *App) renderExpandedToolDisplay(width int, display *toolDisplay, style } func (app *App) toolArgumentLines(width int, display *toolDisplay, style tcell.Style) []tui.Line { + if source := executeSource(display); source != "" { + return app.toolCodeLines(width, source, style) + } + arguments := prettyJSON(display.ArgumentsJSON) if arguments == "" { return nil @@ -115,6 +129,36 @@ func (app *App) toolArgumentLines(width int, display *toolDisplay, style tcell.S return plainSectionLines(width, "args", arguments, style) } +func (app *App) toolCodeLines(width int, source string, style tcell.Style) []tui.Line { + view := tui.CodeBlock{ + Style: style, + Engine: &app.renderer.Lexer, + Language: "go", + Text: source, + Theme: codeTheme(app.theme), + } + content := tui.PrefixLines( + view.Render(max(1, toolContentWidth(width)), maxToolBlockRenderLines), + " ", + style, + ) + + return app.styledSectionLines(width, "code", content, style) +} + +func executeSource(display *toolDisplay) string { + if display == nil { + return "" + } + + name := strings.TrimSpace(display.Name) + if name != toolDisplayExecute && name != toolDisplayWorkflow { + return "" + } + + return rawTextArg(decodeToolArgs(display.ArgumentsJSON), "source") +} + func (app *App) toolDiffLines(width int, display *toolDisplay, style tcell.Style) []tui.Line { diff := toolDisplayDiff(display) if diff == "" { @@ -164,7 +208,12 @@ func toolBlockEnd(width int, style tcell.Style) []tui.Line { } func (app *App) toolHeaderDisplayLines(width int, display *toolDisplay, style tcell.Style) []tui.Line { - return paddedContentLines(width, app.toolDisplayTitle(display), style.Bold(true), false) + title := app.toolDisplayTitle(display) + if display.Nested { + title = " ↳ " + title + } + + return paddedContentLines(width, title, style.Bold(true), false) } func (app *App) toolExpandHintLines(width int, style tcell.Style) []tui.Line { @@ -187,7 +236,13 @@ func plainSectionLines(width int, label, content string, style tcell.Style) []tu func padLinesRight(lines []tui.Line, width int) []tui.Line { padded := make([]tui.Line, 0, len(lines)) for _, line := range lines { - line.Text = tui.PadRight(line.Text, width) + padding := max(0, width-line.Width()) + + line.Text += strings.Repeat(" ", padding) + if len(line.Spans) > 0 && padding > 0 { + line.Spans = append(line.Spans, tui.Span{Style: line.Style, Text: strings.Repeat(" ", padding)}) + } + padded = append(padded, line) } @@ -440,8 +495,10 @@ func prettyJSON(value string) string { func toolDisplayOutput(display *toolDisplay) string { output := strings.Trim(display.Output, "\n") - if display.Error != "" { - output = strings.Trim(display.Error+"\n"+output, "\n") + + errorText := strings.Trim(display.Error, "\n") + if errorText != "" && errorText != output { + output = strings.Trim(errorText+"\n"+output, "\n") } return output @@ -450,6 +507,7 @@ func toolDisplayOutput(display *toolDisplay) string { func parseToolEventContent(content, fallback string) parsedToolEvent { event := parsedToolEvent{ Name: fallback, + ParentCallID: "", ArgumentsJSON: "", DetailsJSON: "", Error: "", @@ -480,6 +538,7 @@ func parseToolEventContent(content, fallback string) parsedToolEvent { } } + event.ParentCallID = strings.TrimSpace(strings.Join(sections[toolSectionParentCallID], "\n")) event.ArgumentsJSON = strings.TrimSpace(strings.Join(sections[toolSectionArguments], "\n")) event.DetailsJSON = strings.TrimSpace(strings.Join(sections[toolSectionDetails], "\n")) event.Error = strings.Trim(strings.Join(sections[toolSectionError], "\n"), "\n") @@ -496,7 +555,8 @@ func parseToolSectionHeader(line string) (name, value string, ok bool) { name = strings.TrimSpace(left) switch name { - case toolSectionTool, toolSectionArguments, toolSectionDetails, toolSectionError, toolSectionOutput: + case toolSectionTool, toolSectionParentCallID, toolSectionArguments, toolSectionDetails, + toolSectionError, toolSectionOutput: return name, strings.TrimSpace(right), true default: return "", "", false diff --git a/internal/terminal/tool_display.go b/internal/terminal/tool_display.go index 57244f17..3386d35c 100644 --- a/internal/terminal/tool_display.go +++ b/internal/terminal/tool_display.go @@ -16,6 +16,8 @@ const ( toolDisplayDefaultPath = "." toolDisplayMaxArgs = 3 toolDisplayGrepOptions = 3 + toolDisplayExecute = "execute" + toolDisplayWorkflow = "workflow" ) type toolSummaryRenderer func(map[string]any, string) string @@ -36,12 +38,14 @@ type toolDisplay struct { Output string Error string Status toolDisplayStatus + Nested bool } -func toolDisplayFromCall(call assistant.ToolCallEvent) toolDisplay { +func toolDisplayFromCall(call *assistant.ToolCallEvent) toolDisplay { return toolDisplay{ Title: toolSummary(call.Name, call.ArgumentsJSON, toolArgumentsMap(call.Arguments)), Name: call.Name, + Nested: call.ParentCallID != "", ArgumentsJSON: call.ArgumentsJSON, DetailsJSON: "", Output: "", @@ -59,6 +63,7 @@ func toolDisplayFromParsedEvent(event *parsedToolEvent) toolDisplay { return toolDisplay{ Title: toolSummary(event.Name, event.ArgumentsJSON, nil), Name: event.Name, + Nested: event.ParentCallID != "", ArgumentsJSON: event.ArgumentsJSON, DetailsJSON: event.DetailsJSON, Output: event.Output, @@ -122,6 +127,10 @@ func otherToolSummaryRenderer(name tool.Name) (toolSummaryRenderer, bool) { return astToolSummary, true case string(tool.NameFetch): return fetchToolSummary, true + case toolDisplayExecute: + return executeToolSummary, true + case toolDisplayWorkflow: + return workflowToolSummary, true default: return nil, false } @@ -144,6 +153,14 @@ func decodeToolArgs(raw string) map[string]any { return args } +func executeToolSummary(_ map[string]any, _ string) string { + return "execute Go" +} + +func workflowToolSummary(_ map[string]any, _ string) string { + return "workflow Go" +} + func bashToolSummary(args map[string]any, fallback string) string { command := stringArg(args, "command") if command == "" { diff --git a/internal/terminal/tool_display_internal_test.go b/internal/terminal/tool_display_internal_test.go index c91d0a3b..d7de858b 100644 --- a/internal/terminal/tool_display_internal_test.go +++ b/internal/terminal/tool_display_internal_test.go @@ -10,6 +10,7 @@ import ( "github.com/omarluq/librecode/internal/assistant" "github.com/omarluq/librecode/internal/testutil" + "github.com/omarluq/librecode/internal/tui" ) func TestToolSummaryHumanizesKnownTools(t *testing.T) { @@ -157,6 +158,7 @@ func TestToolDisplayAdditionalFallbackBranches(t *testing.T) { errorDisplay := toolDisplayFromParsedEvent(&parsedToolEvent{ Name: testToolBash, + ParentCallID: "", ArgumentsJSON: `{"command":"false"}`, DetailsJSON: "", Output: "", @@ -189,11 +191,13 @@ func TestToolSummaryNumericArgumentTypes(t *testing.T) { func TestToolDisplayFromCallUsesStructuredArguments(t *testing.T) { t.Parallel() - display := toolDisplayFromCall(assistant.ToolCallEvent{ + display := toolDisplayFromCall(&assistant.ToolCallEvent{ Arguments: testutil.ToolArguments(map[string]any{testToolCommandKey: "go test ./..."}), ID: "call_1", + ParentCallID: "", Name: testToolBash, ArgumentsJSON: `{"command":"stale"}`, + Sequence: 0, }) assert.Equal(t, "$ go test ./...", display.Title) @@ -281,6 +285,66 @@ func TestTailExpandedToolLinesIncludesErrorWithOutput(t *testing.T) { assert.NotEqual(t, -1, lineIndexContaining(lines, "stdout")) } +func TestToolDisplayOutputDoesNotDuplicateErrorAsOutput(t *testing.T) { + t.Parallel() + + const compileError = "compile failed" + + display := testToolDisplay("execute", toolDisplayError) + display.Error = compileError + display.Output = compileError + + assert.Equal(t, compileError, toolDisplayOutput(&display)) +} + +func TestRenderExpandedExecuteToolHighlightsGoSource(t *testing.T) { + t.Parallel() + + app := newRenderTestApp(t) + app.toolsExpanded = true + display := testToolDisplay("execute", toolDisplaySuccess) + display.Name = toolDisplayExecute + display.ArgumentsJSON = `{"source":"import \"tools\"\nfunc main() { tools.Call(\"read\", map[string]any{}) }"}` + + lines := app.renderToolDisplayBlock(80, &display) + + assertExecuteSourceHighlighted(t, lines) +} + +func TestRenderPendingExpandedExecuteToolHighlightsGoSource(t *testing.T) { + t.Parallel() + + app := newRenderTestApp(t) + app.toolsExpanded = true + display := testToolDisplay("execute Go", toolDisplayPending) + display.Name = toolDisplayExecute + display.ArgumentsJSON = `{"source":"import \"tools\"\nfunc main() { tools.Call(\"read\", map[string]any{}) }"}` + + lines := app.renderToolDisplayBlock(80, &display) + + assertExecuteSourceHighlighted(t, lines) + assert.NotEqual(t, -1, lineIndexContaining(lines, pendingToolIndicator+" execute Go")) +} + +func assertExecuteSourceHighlighted(t *testing.T, lines []tui.Line) { + t.Helper() + + assert.NotEqual(t, -1, lineIndexContaining(lines, "code:")) + assert.NotEqual(t, -1, lineIndexContaining(lines, `import "tools"`)) + assert.Equal(t, -1, lineIndexContaining(lines, `"source":`)) + assert.True(t, hasStyledToolLine(lines, "tools.Call"), "execute source should retain syntax-highlighted spans") +} + +func hasStyledToolLine(lines []tui.Line, text string) bool { + for _, line := range lines { + if strings.Contains(line.Text, text) && len(line.Spans) > 1 { + return true + } + } + + return false +} + func TestRenderExpandedToolDisplayPrettyPrintsArguments(t *testing.T) { t.Parallel() @@ -334,6 +398,7 @@ func testToolDisplay(title string, status toolDisplayStatus) toolDisplay { Output: "", Error: "", Status: status, + Nested: false, } } diff --git a/internal/terminal/workflow_summary_internal_test.go b/internal/terminal/workflow_summary_internal_test.go new file mode 100644 index 00000000..f259d23c --- /dev/null +++ b/internal/terminal/workflow_summary_internal_test.go @@ -0,0 +1,268 @@ +package terminal + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/omarluq/librecode/internal/assistant" + "github.com/omarluq/librecode/internal/database" + "github.com/omarluq/librecode/internal/transcript" +) + +type activeWorkflowInspectorStub struct { + byID map[string]*database.WorkflowRunEntity + runs []database.WorkflowRunEntity +} + +func (stub *activeWorkflowInspectorStub) Get( + _ context.Context, runID string, +) (*database.WorkflowRunEntity, bool, error) { + run, found := stub.byID[runID] + + return run, found, nil +} + +func (stub *activeWorkflowInspectorStub) List( + context.Context, string, int, +) ([]database.WorkflowRunEntity, error) { + return stub.runs, nil +} + +func (stub *activeWorkflowInspectorStub) Events( + context.Context, string, int64, int, +) ([]database.TaskEventEntity, error) { + return nil, nil +} + +func (stub *activeWorkflowInspectorStub) AgentTasks( + context.Context, string, +) ([]database.WorkflowAgentTaskEntity, error) { + return nil, nil +} + +func TestActiveWorkflowAppearsInAgentSummary(t *testing.T) { + t.Parallel() + + app := newRenderTestApp(t) + app.sessionID = workflowTestSessionID + app.workflows = &activeWorkflowInspectorStub{byID: nil, runs: []database.WorkflowRunEntity{ + workflowSummaryRun("active", database.TaskRunning), + workflowSummaryRun(statusDone, database.TaskSucceeded), + }} + + app.refreshActiveWorkflows(t.Context()) + require.Len(t, app.activeWorkflows, 1) + lines := app.renderAgentTaskSummary(80) + require.NotEmpty(t, lines) + assert.Contains(t, lines[0].Text, "workflow(active review)") + assert.True(t, app.hasRunningAgentTasks()) +} + +func TestWorkflowFailureIsPushedIntoCompletedTurn(t *testing.T) { + t.Parallel() + + running := workflowSummaryRun("failed-run", database.TaskRunning) + failed := workflowSummaryRun("failed-run", database.TaskFailed) + failed.Task.ErrorMessage = "compile failed" + stub := &activeWorkflowInspectorStub{ + byID: map[string]*database.WorkflowRunEntity{failed.Task.ID: &failed}, + runs: []database.WorkflowRunEntity{failed}, + } + app := newRenderTestApp(t) + app.sessionID = workflowTestSessionID + app.workflows = stub + app.working = true + app.activeWorkflows = []database.WorkflowRunEntity{running} + + app.refreshActiveWorkflows(t.Context()) + + assert.Empty(t, app.activeWorkflows) + require.Len(t, app.liveAgentCompletions, 1) + assert.Equal(t, transcript.RoleToolResult, app.liveAgentCompletions[0].Role) + assert.Contains(t, app.liveAgentCompletions[0].Content, "active review") + assert.Contains(t, app.liveAgentCompletions[0].Content, "failed-run") + assert.Contains(t, app.liveAgentCompletions[0].Content, "compile failed") + require.Len(t, app.hiddenQueuedMessages, 1) + assert.Contains(t, app.hiddenQueuedMessages[0], "background workflow failed") + + app.refreshActiveWorkflows(t.Context()) + assert.Len(t, app.liveAgentCompletions, 1, "failure must only be delivered once") +} + +func TestTerminalWorkflowWithoutFailureDoesNotPushCompletion(t *testing.T) { + t.Parallel() + + for _, state := range []database.TaskState{database.TaskSucceeded, database.TaskCanceled} { + t.Run(string(state), func(t *testing.T) { + t.Parallel() + + running := workflowSummaryRun("terminal-run", database.TaskRunning) + terminal := workflowSummaryRun("terminal-run", state) + app := newRenderTestApp(t) + app.sessionID = workflowTestSessionID + app.workflows = &activeWorkflowInspectorStub{ + byID: map[string]*database.WorkflowRunEntity{terminal.Task.ID: &terminal}, + runs: []database.WorkflowRunEntity{terminal}, + } + app.activeWorkflows = []database.WorkflowRunEntity{running} + + app.refreshActiveWorkflows(t.Context()) + + assert.Empty(t, app.activeWorkflows) + assert.Empty(t, app.liveAgentCompletions) + assert.Empty(t, app.hiddenQueuedMessages) + }) + } +} + +func TestTrackStartedWorkflowDeliversImmediateFailureOnce(t *testing.T) { + t.Parallel() + + failed := workflowSummaryRun("immediate-failure", database.TaskFailed) + failed.Task.ErrorMessage = "source did not compile" + app := newRenderTestApp(t) + app.sessionID = workflowTestSessionID + app.working = true + app.workflows = &activeWorkflowInspectorStub{ + byID: map[string]*database.WorkflowRunEntity{failed.Task.ID: &failed}, + runs: nil, + } + event := &assistant.ToolEvent{ + CallID: "", ParentCallID: "", Sequence: 0, Name: workflowToolName, + ArgumentsJSON: "", DetailsJSON: `{"run_id":"immediate-failure"}`, + Result: "", Error: "", IsError: false, + } + + app.trackStartedWorkflow(t.Context(), event) + app.trackStartedWorkflow(t.Context(), event) + + assert.Empty(t, app.activeWorkflows) + require.Len(t, app.liveAgentCompletions, 1) + assert.Equal(t, transcript.RoleToolResult, app.liveAgentCompletions[0].Role) + assert.Contains(t, app.liveAgentCompletions[0].Content, "active review") + assert.Contains(t, app.liveAgentCompletions[0].Content, failed.Task.ID) + assert.Contains(t, app.liveAgentCompletions[0].Content, failed.Task.ErrorMessage) + require.Len(t, app.hiddenQueuedMessages, 1) +} + +func TestWorkflowFailureNotificationFallbacksAndBusyBranches(t *testing.T) { + t.Parallel() + + tests := []struct { + makeBusy func(*App) + name string + }{ + {name: workflowTestWorking, makeBusy: func(app *App) { app.working = true }}, + {name: "auth working", makeBusy: func(app *App) { app.authWorking = true }}, + {name: workflowTestCompacting, makeBusy: func(app *App) { app.compacting = true }}, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + failed := workflowSummaryRun("fallback-failure", database.TaskFailed) + failed.Name = " " + failed.Task.ErrorMessage = " " + app := newRenderTestApp(t) + testCase.makeBusy(app) + + app.deliverWorkflowFailure(t.Context(), nil) + + nonfailure := failed + nonfailure.Task.State = database.TaskSucceeded + app.deliverWorkflowFailure(t.Context(), &nonfailure) + app.deliverWorkflowFailure(t.Context(), &failed) + + assert.Equal(t, "workflow failed", app.statusMessage) + require.Len(t, app.liveAgentCompletions, 1) + assert.Contains(t, app.liveAgentCompletions[0].Content, toolDisplayWorkflow) + assert.Contains(t, app.liveAgentCompletions[0].Content, "No error detail was returned.") + require.Len(t, app.hiddenQueuedMessages, 1) + assert.Contains(t, app.hiddenQueuedMessages[0], "background workflow failed") + }) + } +} + +func TestTrackStartedWorkflowRejectsForeignRunAndInspectorError(t *testing.T) { + t.Parallel() + + foreign := workflowSummaryRun("foreign-run", database.TaskRunning) + foreign.Task.OwnerSessionID = workflowTestForeignSession + + tests := []struct { + inspector *workflowInspectorStub + name string + }{ + { + name: "foreign session", + inspector: &workflowInspectorStub{ + listErr: nil, getErr: nil, eventsErr: nil, agentTasksErr: nil, + getRun: &foreign, runs: nil, events: nil, children: nil, found: true, + }, + }, + { + name: "get error", + inspector: &workflowInspectorStub{ + listErr: nil, getErr: assert.AnError, eventsErr: nil, agentTasksErr: nil, + getRun: nil, runs: nil, events: nil, children: nil, found: false, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + app := newRenderTestApp(t) + app.sessionID = workflowTestSessionID + app.workflows = testCase.inspector + app.trackStartedWorkflow(t.Context(), &assistant.ToolEvent{ + CallID: "", ParentCallID: "", Sequence: 0, Name: workflowToolName, + ArgumentsJSON: "", DetailsJSON: `{"run_id":"foreign-run"}`, + Result: "", Error: "", IsError: false, + }) + + assert.Empty(t, app.activeWorkflows) + assert.Empty(t, app.liveAgentCompletions) + }) + } +} + +func TestTrackStartedWorkflowUsesToolResultRunID(t *testing.T) { + t.Parallel() + + run := workflowSummaryRun("queued-run", database.TaskQueued) + app := newRenderTestApp(t) + app.sessionID = workflowTestSessionID + app.workflows = &activeWorkflowInspectorStub{ + byID: map[string]*database.WorkflowRunEntity{run.Task.ID: &run}, runs: nil, + } + + app.trackStartedWorkflow(t.Context(), &assistant.ToolEvent{ + CallID: "", ParentCallID: "", Sequence: 0, Name: workflowToolName, + ArgumentsJSON: "", DetailsJSON: `{"run_id":" queued-run "}`, + Result: "", Error: "", IsError: false, + }) + + require.Len(t, app.activeWorkflows, 1) + assert.Equal(t, run.Task.ID, app.activeWorkflows[0].Task.ID) + assert.Equal(t, "queued-run", workflowRunIDFromDetails(`{"run_id":" queued-run "}`)) + assert.Empty(t, workflowRunIDFromDetails(`{`)) +} + +func workflowSummaryRun(id string, state database.TaskState) database.WorkflowRunEntity { + return database.WorkflowRunEntity{ + Task: database.TaskEntity{ + CreatedAt: time.Time{}, StartedAt: nil, FinishedAt: nil, UpdatedAt: time.Time{}, + LeaseExpiresAt: nil, ID: id, Kind: database.TaskKindWorkflow, ParentTaskID: "", + OwnerSessionID: workflowTestSessionID, ConcurrencyKey: "", LeaseOwner: "", State: state, + Result: "", ErrorCode: "", ErrorMessage: "", + }, + Name: "active review", Source: "", SourceHash: "", SourceVersion: "", ArgumentsJSON: "", + } +} diff --git a/internal/terminal/workflow_test_helpers_internal_test.go b/internal/terminal/workflow_test_helpers_internal_test.go new file mode 100644 index 00000000..07a9949e --- /dev/null +++ b/internal/terminal/workflow_test_helpers_internal_test.go @@ -0,0 +1,41 @@ +package terminal + +import ( + "context" + + "github.com/omarluq/librecode/internal/database" +) + +const ( + workflowTestSessionID = "session-1" + workflowTestForeignSession = "another-session" + workflowTestWorking = "working" + workflowTestCompacting = "compacting" +) + +type workflowInspectorStub struct { + listErr error + getErr error + eventsErr error + agentTasksErr error + getRun *database.WorkflowRunEntity + runs []database.WorkflowRunEntity + events []database.TaskEventEntity + children []database.WorkflowAgentTaskEntity + found bool +} + +func (stub *workflowInspectorStub) Get( + context.Context, + string, +) (*database.WorkflowRunEntity, bool, error) { + return stub.getRun, stub.found, stub.getErr +} + +func (stub *workflowInspectorStub) List( + context.Context, + string, + int, +) ([]database.WorkflowRunEntity, error) { + return stub.runs, stub.listErr +} diff --git a/internal/tool/registry.go b/internal/tool/registry.go index 7585e2b9..e0ed25db 100644 --- a/internal/tool/registry.go +++ b/internal/tool/registry.go @@ -116,6 +116,17 @@ func (registry *Registry) Definitions() []Definition { return definitions } +// Has reports whether a tool is registered under name. +func (registry *Registry) Has(name Name) bool { + if registry == nil { + return false + } + + _, exists := registry.executors[name] + + return exists +} + // Execute runs a named tool with raw JSON object arguments. func (registry *Registry) Execute(ctx context.Context, name string, input Arguments) (Result, error) { executor, ok := registry.executors[Name(name)] diff --git a/internal/tool/registry_test.go b/internal/tool/registry_test.go index d1491834..cc8c62ac 100644 --- a/internal/tool/registry_test.go +++ b/internal/tool/registry_test.go @@ -166,6 +166,9 @@ func TestRegistry_Metadata(t *testing.T) { registry := tool.NewRegistry(cwd) assert.Equal(t, cwd, registry.CWD()) + assert.True(t, registry.Has(tool.NameRead)) + assert.False(t, registry.Has("missing")) + assert.False(t, (*tool.Registry)(nil).Has(tool.NameRead)) assert.NotEmpty(t, registry.Definitions()) assert.Len(t, tool.AllDefinitions(), len(registry.Definitions())) } diff --git a/internal/transcript/tool_event.go b/internal/transcript/tool_event.go index 470bcf9a..d7ae8c65 100644 --- a/internal/transcript/tool_event.go +++ b/internal/transcript/tool_event.go @@ -2,15 +2,21 @@ // assistant persistence and terminal presentation. package transcript -import "strings" +import ( + "strconv" + "strings" +) // ToolEvent captures the display/persistence fields of a completed tool call. type ToolEvent struct { + CallID string + ParentCallID string Name string ArgumentsJSON string DetailsJSON string Result string Error string + Sequence int IsError bool } @@ -29,7 +35,7 @@ func formatToolEvent(event *ToolEvent, includeStructuredError bool) string { return "tool: " } - parts := []string{"tool: " + event.Name} + parts := toolIdentityParts(event) if strings.TrimSpace(event.ArgumentsJSON) != "" { parts = append(parts, "arguments:", event.ArgumentsJSON) } @@ -52,3 +58,20 @@ func formatToolEvent(event *ToolEvent, includeStructuredError bool) string { return strings.Join(parts, "\n") } + +func toolIdentityParts(event *ToolEvent) []string { + parts := []string{"tool: " + event.Name} + if event.CallID != "" { + parts = append(parts, "call_id: "+event.CallID) + } + + if event.ParentCallID != "" { + parts = append(parts, "parent_call_id: "+event.ParentCallID) + } + + if event.Sequence > 0 { + parts = append(parts, "sequence: "+strconv.Itoa(event.Sequence)) + } + + return parts +} diff --git a/internal/transcript/tool_event_test.go b/internal/transcript/tool_event_test.go index 044b2b32..bea9f8c5 100644 --- a/internal/transcript/tool_event_test.go +++ b/internal/transcript/tool_event_test.go @@ -13,11 +13,14 @@ func TestFormatToolEventPersistence(t *testing.T) { t.Parallel() event := transcript.ToolEvent{ + CallID: "", + ParentCallID: "", Name: "bash", ArgumentsJSON: `{"command":"false"}`, DetailsJSON: `{"exit_code":1}`, Result: "stderr output", Error: "exit status 1", + Sequence: 0, IsError: true, } @@ -35,15 +38,32 @@ func TestFormatToolEventPersistence(t *testing.T) { ), transcript.FormatToolEventPersistence(&event)) } +func TestFormatToolEventPersistenceIncludesTraceIdentity(t *testing.T) { + t.Parallel() + + event := transcript.ToolEvent{ + CallID: "outer/2", ParentCallID: "outer", Name: "read", ArgumentsJSON: "", DetailsJSON: "", + Result: "ok", Error: "", Sequence: 2, IsError: false, + } + + formatted := transcript.FormatToolEventPersistence(&event) + assert.Contains(t, formatted, "call_id: outer/2") + assert.Contains(t, formatted, "parent_call_id: outer") + assert.Contains(t, formatted, "sequence: 2") +} + func TestFormatToolEventDisplayOmitsStructuredErrorMarker(t *testing.T) { t.Parallel() event := transcript.ToolEvent{ + CallID: "", + ParentCallID: "", Name: "read", ArgumentsJSON: "", DetailsJSON: "", Result: "", Error: "read failed", + Sequence: 0, IsError: true, } @@ -58,11 +78,14 @@ func TestFormatToolEventSkipsBlankOptionalSections(t *testing.T) { t.Parallel() event := transcript.ToolEvent{ + CallID: "", + ParentCallID: "", Name: "write", ArgumentsJSON: " \n\t ", DetailsJSON: "", Result: "\n", Error: "", + Sequence: 0, IsError: false, } diff --git a/internal/workflow/dispatcher.go b/internal/workflow/dispatcher.go new file mode 100644 index 00000000..9fa8133f --- /dev/null +++ b/internal/workflow/dispatcher.go @@ -0,0 +1,210 @@ +package workflow + +import ( + "context" + "errors" + "log/slog" + "sync" + "time" + + "github.com/samber/oops" + + "github.com/omarluq/librecode/internal/database" +) + +const ( + defaultWorkerConcurrency = 2 + defaultDispatchInterval = time.Second + defaultDispatchBuffer = 128 +) + +// DispatcherOptions configures durable workflow queue polling and execution. +type DispatcherOptions struct { + Service *Service + Tasks *database.TaskRepository + Logger *slog.Logger + Concurrency int + Buffer int + Interval time.Duration +} + +// Dispatcher polls durable queued workflows and executes atomically claimed runs. +type Dispatcher struct { + service *Service + tasks *database.TaskRepository + logger *slog.Logger + queue chan string + cancel context.CancelFunc + done <-chan struct{} + wg sync.WaitGroup + submits sync.WaitGroup + interval time.Duration + mu sync.Mutex + closed bool +} + +// NewDispatcher starts durable workflow workers. +func NewDispatcher(ctx context.Context, options DispatcherOptions) (*Dispatcher, error) { + if ctx == nil || options.Service == nil || options.Tasks == nil { + return nil, oops.In("workflow").Code("invalid_dispatcher_dependencies"). + Errorf("process context, workflow service, and task repository are required") + } + + concurrency := options.Concurrency + if concurrency <= 0 { + concurrency = defaultWorkerConcurrency + } + + buffer := options.Buffer + if buffer <= 0 { + buffer = defaultDispatchBuffer + } + + interval := options.Interval + if interval <= 0 { + interval = defaultDispatchInterval + } + + logger := options.Logger + if logger == nil { + logger = slog.Default() + } + + dispatchCtx, cancel := context.WithCancel(ctx) + dispatcher := &Dispatcher{ + service: options.Service, tasks: options.Tasks, logger: logger, + queue: make(chan string, buffer), cancel: cancel, done: dispatchCtx.Done(), + wg: sync.WaitGroup{}, submits: sync.WaitGroup{}, interval: interval, mu: sync.Mutex{}, closed: false, + } + + for range concurrency { + dispatcher.wg.Add(1) + go dispatcher.worker(dispatchCtx) + } + + dispatcher.wg.Add(1) + go dispatcher.poll(dispatchCtx) + + return dispatcher, nil +} + +// Submit persists and schedules a workflow run. +func (dispatcher *Dispatcher) Submit( + ctx context.Context, + request *ServiceRequest, +) (*database.WorkflowRunEntity, error) { + dispatcher.mu.Lock() + if dispatcher.closed { + dispatcher.mu.Unlock() + + return nil, oops.In("workflow").Code("dispatcher_shutdown").Errorf("workflow dispatcher is shut down") + } + + dispatcher.submits.Add(1) + dispatcher.mu.Unlock() + + defer dispatcher.submits.Done() + + run, err := dispatcher.service.Submit(ctx, request) + if err != nil { + return nil, err + } + + dispatcher.enqueue(run.Task.ID) + + return run, nil +} + +// Await delegates durable completion waiting to the workflow service. +func (dispatcher *Dispatcher) Await(ctx context.Context, runID string) (*database.WorkflowRunEntity, error) { + return dispatcher.service.Await(ctx, runID) +} + +// Shutdown stops polling and waits for workers. +func (dispatcher *Dispatcher) Shutdown(ctx context.Context) error { + dispatcher.mu.Lock() + dispatcher.closed = true + dispatcher.cancel() + dispatcher.mu.Unlock() + + done := make(chan struct{}) + + go func() { + dispatcher.submits.Wait() + dispatcher.wg.Wait() + close(done) + }() + + select { + case <-done: + return nil + case <-ctx.Done(): + return oops.In("workflow").Code("dispatcher_shutdown").Wrapf(ctx.Err(), "shutdown workflow dispatcher") + } +} + +func (dispatcher *Dispatcher) poll(ctx context.Context) { + defer dispatcher.wg.Done() + + dispatcher.enqueueQueued(ctx) + + ticker := time.NewTicker(dispatcher.interval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + dispatcher.enqueueQueued(ctx) + } + } +} + +func (dispatcher *Dispatcher) enqueueQueued(ctx context.Context) { + tasks, err := dispatcher.tasks.ListByStates( + ctx, + database.TaskKindWorkflow, + []database.TaskState{database.TaskQueued}, + cap(dispatcher.queue), + ) + if err != nil { + if !errors.Is(err, context.Canceled) { + dispatcher.logger.Error("list queued workflows", "error", err) + } + + return + } + + for index := range tasks { + dispatcher.enqueue(tasks[index].ID) + } +} + +func (dispatcher *Dispatcher) enqueue(runID string) { + select { + case dispatcher.queue <- runID: + case <-dispatcher.done: + default: + } +} + +func (dispatcher *Dispatcher) worker(ctx context.Context) { + defer dispatcher.wg.Done() + + for { + select { + case <-ctx.Done(): + return + case runID := <-dispatcher.queue: + if ctx.Err() != nil { + return + } + + if _, err := dispatcher.service.ExecuteQueued(ctx, runID); err != nil && + !errors.Is(err, context.Canceled) { + dispatcher.logger.Error("execute queued workflow", "run_id", runID, "error", err) + } + } + } +} diff --git a/internal/workflow/dispatcher_cross_process_test.go b/internal/workflow/dispatcher_cross_process_test.go new file mode 100644 index 00000000..65db756e --- /dev/null +++ b/internal/workflow/dispatcher_cross_process_test.go @@ -0,0 +1,66 @@ +package workflow_test + +import ( + "context" + "database/sql" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + _ "modernc.org/sqlite" // Register the sqlite database/sql driver used by this cross-process test. + + "github.com/omarluq/librecode/internal/database" + "github.com/omarluq/librecode/internal/workflow" +) + +func TestDispatcherClaimsWorkflowSubmittedByAnotherService(t *testing.T) { + t.Parallel() + + databasePath := t.TempDir() + "/workflow.db" + dsn := database.SQLiteDSN(databasePath, database.SQLiteOptions{BusyTimeout: time.Second}) + submitterDB, err := sql.Open("sqlite", dsn) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, submitterDB.Close()) }) + require.NoError(t, database.Migrate(t.Context(), submitterDB)) + + workerDB, err := sql.Open("sqlite", dsn) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, workerDB.Close()) }) + + repository := database.NewWorkflowRepository(submitterDB) + sessions := database.NewSessionRepository(submitterDB) + owner, err := sessions.CreateSession(t.Context(), t.TempDir(), "cross-process workflow", "") + require.NoError(t, err) + + submitterRunner, err := workflow.NewRunner(newFakeController()) + require.NoError(t, err) + submitter, err := workflow.NewService(repository, submitterRunner) + require.NoError(t, err) + + workerRunner, err := workflow.NewRunner(newFakeController()) + require.NoError(t, err) + worker, err := workflow.NewService(database.NewWorkflowRepository(workerDB), workerRunner) + require.NoError(t, err) + + run, err := submitter.Submit(t.Context(), &workflow.ServiceRequest{ + Name: "cross-process", Source: "21 * 2", SourceVersion: "v1", ArgumentsJSON: "{}", + OwnerSessionID: owner.ID, + }) + require.NoError(t, err) + + dispatcher, err := workflow.NewDispatcher(context.Background(), workflow.DispatcherOptions{ + Service: worker, + Tasks: database.NewTaskRepository(workerDB), + Logger: nil, + Concurrency: 1, + Buffer: 4, + Interval: 10 * time.Millisecond, + }) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, dispatcher.Shutdown(context.Background())) }) + + completed, err := submitter.Await(t.Context(), run.Task.ID) + require.NoError(t, err) + assert.Equal(t, database.TaskSucceeded, completed.Task.State) +} diff --git a/internal/workflow/dispatcher_test.go b/internal/workflow/dispatcher_test.go new file mode 100644 index 00000000..ad0c1498 --- /dev/null +++ b/internal/workflow/dispatcher_test.go @@ -0,0 +1,172 @@ +package workflow_test + +import ( + "context" + "database/sql" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/omarluq/librecode/internal/database" + "github.com/omarluq/librecode/internal/workflow" +) + +const dispatcherQueuedName = "queued" + +func TestDispatcherDependencyDefaultsAndShutdownBoundaries(t *testing.T) { + t.Parallel() + + service, repository, _ := newWorkflowService(t, newFakeController()) + + tests := []struct { + ctx func() context.Context + service *workflow.Service + tasks *database.TaskRepository + name string + }{ + {name: "nil context", ctx: func() context.Context { return nil }, service: service, tasks: repository.Tasks()}, + {name: "nil service", ctx: t.Context, service: nil, tasks: repository.Tasks()}, + {name: "nil tasks", ctx: t.Context, service: service, tasks: nil}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + dispatcher, err := workflow.NewDispatcher(test.ctx(), workflow.DispatcherOptions{ + Service: test.service, + Tasks: test.tasks, + Logger: nil, + Concurrency: 0, + Buffer: 0, + Interval: 0, + }) + require.ErrorContains(t, err, "are required") + assert.Nil(t, dispatcher) + }) + } + + dispatcher, err := workflow.NewDispatcher(t.Context(), workflow.DispatcherOptions{ + Service: service, + Tasks: repository.Tasks(), + Logger: nil, + Concurrency: 0, + Buffer: 0, + Interval: 0, + }) + require.NoError(t, err) + + _, err = dispatcher.Submit(t.Context(), nil) + require.ErrorContains(t, err, "request is required") + require.NoError(t, dispatcher.Shutdown(t.Context())) + + run, err := dispatcher.Submit(t.Context(), &workflow.ServiceRequest{ + Name: "", + Source: "", + SourceVersion: "", + ArgumentsJSON: "", + OwnerSessionID: "", + }) + require.ErrorContains(t, err, "dispatcher is shut down") + assert.Nil(t, run) +} + +func TestDispatcherShutdownDoesNotHoldLifecycleLockAcrossSubmitIO(t *testing.T) { + t.Parallel() + + connection, err := sql.Open("sqlite", "file:"+t.Name()+"?mode=memory&cache=shared&_pragma=foreign_keys(1)") + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, connection.Close()) }) + connection.SetMaxOpenConns(1) + require.NoError(t, database.Migrate(t.Context(), connection)) + + repository := database.NewWorkflowRepository(connection) + sessions := database.NewSessionRepository(connection) + owner, err := sessions.CreateSession(t.Context(), "/work", "owner", "") + require.NoError(t, err) + runner, err := workflow.NewRunner(newFakeController()) + require.NoError(t, err) + service, err := workflow.NewService(repository, runner) + require.NoError(t, err) + transaction, err := connection.BeginTx(t.Context(), nil) + require.NoError(t, err) + + initialWaitCount := connection.Stats().WaitCount + + dispatcher, err := workflow.NewDispatcher(context.Background(), workflow.DispatcherOptions{ + Service: service, Tasks: repository.Tasks(), Logger: nil, Concurrency: 1, Buffer: 1, Interval: time.Hour, + }) + require.NoError(t, err) + + submitDone := make(chan error, 1) + + go func() { + _, submitErr := dispatcher.Submit(context.Background(), &workflow.ServiceRequest{ + Name: "blocked", Source: "1", SourceVersion: "v1", ArgumentsJSON: "{}", OwnerSessionID: owner.ID, + }) + submitDone <- submitErr + }() + + require.Eventually(t, func() bool { + // The dispatcher poller and Submit must both be waiting behind the transaction. + return connection.Stats().WaitCount >= initialWaitCount+2 + }, 5*time.Second, time.Millisecond, "Submit did not block waiting for the database connection") + + shutdownCtx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + + require.ErrorIs(t, dispatcher.Shutdown(shutdownCtx), context.DeadlineExceeded) + + _, err = dispatcher.Submit(t.Context(), &workflow.ServiceRequest{ + Name: "", Source: "", SourceVersion: "", ArgumentsJSON: "", OwnerSessionID: "", + }) + require.ErrorContains(t, err, "dispatcher is shut down") + require.NoError(t, transaction.Rollback()) + require.NoError(t, <-submitDone) + require.NoError(t, dispatcher.Shutdown(context.Background())) +} + +func TestDispatcherExecutesSubmittedWorkflow(t *testing.T) { + t.Parallel() + + service, repository, owner := newWorkflowService(t, newFakeController()) + dispatcher, err := workflow.NewDispatcher(context.Background(), workflow.DispatcherOptions{ + Service: service, Tasks: repository.Tasks(), Logger: nil, + Concurrency: 1, Buffer: 4, Interval: 10 * time.Millisecond, + }) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, dispatcher.Shutdown(context.Background())) }) + + run, err := dispatcher.Submit(t.Context(), &workflow.ServiceRequest{ + Name: dispatcherQueuedName, Source: "1 + 1", SourceVersion: "v1", ArgumentsJSON: "{}", + OwnerSessionID: owner, + }) + require.NoError(t, err) + + completed, err := dispatcher.Await(t.Context(), run.Task.ID) + require.NoError(t, err) + assert.Equal(t, database.TaskSucceeded, completed.Task.State) +} + +func TestDispatcherRecoversQueuedWorkflow(t *testing.T) { + t.Parallel() + + service, repository, owner := newWorkflowService(t, newFakeController()) + run, err := service.Submit(t.Context(), &workflow.ServiceRequest{ + Name: "recovered", Source: "2 + 2", SourceVersion: "v1", ArgumentsJSON: "{}", + OwnerSessionID: owner, + }) + require.NoError(t, err) + + dispatcher, err := workflow.NewDispatcher(context.Background(), workflow.DispatcherOptions{ + Service: service, Tasks: repository.Tasks(), Logger: nil, + Concurrency: 1, Buffer: 4, Interval: 10 * time.Millisecond, + }) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, dispatcher.Shutdown(context.Background())) }) + + completed, err := dispatcher.Await(t.Context(), run.Task.ID) + require.NoError(t, err) + assert.Equal(t, database.TaskSucceeded, completed.Task.State) +} diff --git a/internal/workflow/pipeline.go b/internal/workflow/pipeline.go new file mode 100644 index 00000000..25f92d59 --- /dev/null +++ b/internal/workflow/pipeline.go @@ -0,0 +1,10 @@ +package workflow + +const pipelineNotScheduled = "pipeline stopped before item was scheduled" + +// PipelineResult is one callback outcome. Results retain input order. +type PipelineResult struct { + Value any + Error string + Index int +} diff --git a/internal/workflow/service.go b/internal/workflow/service.go new file mode 100644 index 00000000..fc621340 --- /dev/null +++ b/internal/workflow/service.go @@ -0,0 +1,635 @@ +package workflow + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "sync" + "time" + + "github.com/gofrs/uuid/v5" + "github.com/samber/oops" + + "github.com/omarluq/librecode/internal/database" +) + +var errRunClaimConflict = errors.New("workflow run is not claimable") + +const ( + defaultSourceVersion = "v1" + workflowStartedEvent = "workflow_started" + workflowResumedEvent = "workflow_resumed" + workflowEventKind = "workflow_event" + workflowSucceeded = "workflow_succeeded" + workflowFailed = "workflow_failed" + workflowCanceled = "workflow_canceled" + workflowInterrupted = "workflow_interrupted" + workflowLeaseDuration = 30 * time.Second + workflowLeaseInterval = 5 * time.Second + workflowAwaitInterval = 100 * time.Millisecond +) + +// ServiceRequest describes one durable one-shot workflow execution. +type ServiceRequest struct { + Name string + Source string + SourceVersion string + ArgumentsJSON string + OwnerSessionID string +} + +type eventPayload struct { + Task TaskResult `json:"task"` + Kind EventKind `json:"kind"` + TaskID string `json:"task_id"` + NodeKey string `json:"node_key,omitempty"` + InvocationIndex int `json:"invocation_index,omitempty"` +} + +type runResultPayload struct { + Value any `json:"value"` + Stdout string `json:"stdout"` + Stderr string `json:"stderr"` + LaunchedTaskIDs []string `json:"launched_task_ids"` + TaskResults []TaskResult `json:"task_results"` +} + +// Service persists workflow lifecycle and delegates source evaluation to a Runner. +type Service struct { + runs *database.WorkflowRepository + runner *Runner + active map[string]context.CancelFunc + leaseOwner string + mu sync.Mutex +} + +// NewService creates the durable workflow execution boundary. +func NewService(runs *database.WorkflowRepository, runner *Runner) (*Service, error) { + if runs == nil || runner == nil { + return nil, oops.In("workflow").Code("invalid_service_dependencies"). + Errorf("workflow repository and runner are required") + } + + leaseOwner, err := uuid.NewV7() + if err != nil { + return nil, oops.In("workflow").Code("create_lease_owner").Wrapf(err, "create workflow lease owner") + } + + return &Service{ + runs: runs, runner: runner, active: make(map[string]context.CancelFunc), + leaseOwner: leaseOwner.String(), mu: sync.Mutex{}, + }, nil +} + +// Submit durably queues a workflow run without executing it. +func (service *Service) Submit( + ctx context.Context, + request *ServiceRequest, +) (*database.WorkflowRunEntity, error) { + return service.createRun(ctx, request) +} + +// Run creates and executes one durable workflow run. +func (service *Service) Run( + ctx context.Context, + request *ServiceRequest, +) (*database.WorkflowRunEntity, *RunResult, error) { + run, persisted, arguments, err := service.prepareRun(ctx, request) + if err != nil { + return run, nil, err + } + + return service.executeExisting(ctx, run, persisted, arguments, request.Name, nil, false) +} + +// ExecuteQueued claims and executes a previously submitted workflow run. +// A false return means another worker claimed the run first. +func (service *Service) ExecuteQueued(ctx context.Context, runID string) (bool, error) { + run, err := service.loadVerifiedRun(ctx, runID) + if err != nil { + return false, err + } + + arguments, err := decodeArguments(run.ArgumentsJSON) + if err != nil { + return false, err + } + + completed, _, err := service.executeExisting(ctx, run, run, arguments, run.Name, nil, false) + if err != nil && errors.Is(err, errRunClaimConflict) { + return false, nil + } + + // Execution failures are durable workflow outcomes, not dispatcher failures. + // Once the terminal state has been persisted, callers can inspect the run + // without the background worker also writing the source error to the TUI. + if err != nil && completed != nil && terminal(completed.Task.State) { + return true, nil + } + + return true, err +} + +// Await waits for a workflow run to reach a terminal state. +func (service *Service) Await(ctx context.Context, runID string) (*database.WorkflowRunEntity, error) { + ticker := time.NewTicker(workflowAwaitInterval) + defer ticker.Stop() + + for { + run, found, err := service.Get(ctx, runID) + if err != nil { + return nil, err + } + + if !found { + return nil, oops.In("workflow").Code("missing_run").Errorf("workflow run was not found") + } + + if terminal(run.Task.State) { + return run, nil + } + + select { + case <-ctx.Done(): + return nil, oops.In("workflow").Code("await_canceled").Wrapf(ctx.Err(), "await workflow run") + case <-ticker.C: + } + } +} + +// Resume replays an interrupted workflow while reusing its persisted child invocations. +func (service *Service) Resume( + ctx context.Context, + runID string, +) (*database.WorkflowRunEntity, *RunResult, error) { + run, err := service.loadVerifiedRun(ctx, runID) + if err != nil { + return nil, nil, err + } + + arguments, err := decodeArguments(run.ArgumentsJSON) + if err != nil { + return run, nil, err + } + + links, err := service.runs.ListAgentTasks(ctx, runID) + if err != nil { + return run, nil, oops.In("workflow").Code("load_agent_links").Wrapf(err, "load workflow agent links") + } + + return service.executeExisting(ctx, run, run, arguments, run.Name, links, true) +} + +func (service *Service) executeExisting( + ctx context.Context, + run *database.WorkflowRunEntity, + persisted *database.WorkflowRunEntity, + arguments map[string]any, + name string, + links []database.WorkflowAgentTaskEntity, + resume bool, +) (*database.WorkflowRunEntity, *RunResult, error) { + claim := &database.TaskClaim{TaskID: run.Task.ID, LeaseOwner: service.leaseOwner, + EventKind: workflowStartedEvent, LeaseExpiresAt: time.Now().UTC().Add(workflowLeaseDuration)} + claimRun := service.runs.Tasks().ClaimQueued + + if resume { + claim.EventKind = workflowResumedEvent + claimRun = service.runs.Tasks().ClaimInterrupted + } + + changed, err := claimRun(ctx, claim) + if err != nil { + return run, nil, oops.In("workflow").Code("start_run").Wrapf(err, "start workflow run") + } + + if !changed { + return run, nil, oops.In("workflow").Code("start_run_conflict"). + Wrapf(errRunClaimConflict, "workflow run is not resumable") + } + + runCtx, cancelRun := context.WithCancel(ctx) + heartbeatCtx, stopHeartbeat := context.WithCancel(runCtx) + + service.track(run.Task.ID, cancelRun) + + defer func() { + stopHeartbeat() + cancelRun() + service.untrack(run.Task.ID) + }() + + heartbeatErr := service.monitorLease(heartbeatCtx, run.Task.ID, cancelRun) + result, runErr := service.runner.Run(runCtx, &RunRequest{ + RunID: persisted.Task.ID, OnEvent: service.eventSink(persisted.Task.ID), Name: name, + Source: persisted.Source, OwnerSessionID: persisted.Task.OwnerSessionID, Arguments: arguments, + PersistedLinks: links, + }) + + stopHeartbeat() + + if heartbeatRunErr := <-heartbeatErr; heartbeatRunErr != nil { + runErr = errors.Join(runErr, heartbeatRunErr) + } + + if finishErr := service.finish(runCtx, run.Task.ID, &result, runErr); finishErr != nil { + return run, &result, errors.Join(runErr, finishErr) + } + + completed, found, err := service.runs.Get(context.WithoutCancel(ctx), run.Task.ID) + if err != nil { + return run, &result, oops.In("workflow").Code("reload_run").Wrapf(err, "reload workflow run") + } + + if !found { + return run, &result, oops.In("workflow").Code("missing_run").Errorf("completed workflow run was not found") + } + + return completed, &result, runErr +} + +func (service *Service) prepareRun( + ctx context.Context, + request *ServiceRequest, +) (run, persisted *database.WorkflowRunEntity, arguments map[string]any, err error) { + if request == nil { + return nil, nil, nil, oops.In("workflow").Code("invalid_service_request").Errorf("request is required") + } + + run, err = service.createRun(ctx, request) + if err != nil { + return nil, nil, nil, err + } + + persisted, err = service.loadVerifiedRun(ctx, run.Task.ID) + if err != nil { + return run, nil, nil, err + } + + arguments, err = decodeArguments(persisted.ArgumentsJSON) + if err != nil { + return run, nil, nil, err + } + + return run, persisted, arguments, nil +} + +// RecoverInterrupted marks abandoned in-process runs interrupted after restart. +// Workflow source is not replayed because interpreter memory is intentionally not persisted. +func (service *Service) RecoverInterrupted(ctx context.Context) ([]string, error) { + runIDs, err := service.runs.Tasks().RecoverExpired(ctx, &database.TaskRecovery{ + ExpiresBefore: time.Now().UTC(), Kind: database.TaskKindWorkflow, + EventKind: workflowInterrupted, ErrorCode: "process_restart", + ErrorMessage: "workflow interrupted by process restart", PayloadJSON: "{}", + TargetState: database.TaskInterrupted, + }) + if err != nil { + return nil, oops.In("workflow").Code("recover_runs").Wrapf(err, "recover interrupted workflow runs") + } + + return runIDs, nil +} + +// Cancel cancels an active run owned by the supplied session. +func (service *Service) Cancel(ctx context.Context, ownerSessionID, runID string) (bool, error) { + run, found, err := service.runs.Get(ctx, runID) + if err != nil { + return false, oops.In("workflow").Code("get_run_for_cancel").Wrapf(err, "get workflow run") + } + + if !found { + return false, nil + } + + if run.Task.OwnerSessionID != ownerSessionID { + return false, nil + } + + changed, err := service.runs.Tasks().Transition( + ctx, runID, []database.TaskState{database.TaskQueued}, database.TaskCanceled, workflowCanceled, + ) + if err != nil { + return false, oops.In("workflow").Code("cancel_queued_run").Wrapf(err, "cancel queued workflow run") + } + + if !changed { + changed, err = service.runs.Tasks().Transition( + ctx, runID, []database.TaskState{database.TaskRunning}, database.TaskCanceling, workflowCanceled, + ) + if err != nil { + return false, oops.In("workflow").Code("cancel_running_run").Wrapf(err, "cancel running workflow run") + } + } + + service.mu.Lock() + cancel := service.active[runID] + service.mu.Unlock() + + if cancel != nil { + cancel() + } + + return changed, nil +} + +// Get returns one workflow run. +func (service *Service) Get(ctx context.Context, runID string) (*database.WorkflowRunEntity, bool, error) { + run, found, err := service.runs.Get(ctx, runID) + if err != nil { + return nil, false, oops.In("workflow").Code("get_run").Wrapf(err, "get workflow run") + } + + return run, found, nil +} + +// List returns workflow runs owned by a session. +func (service *Service) List( + ctx context.Context, + ownerSessionID string, + limit int, +) ([]database.WorkflowRunEntity, error) { + runs, err := service.runs.ListByOwner(ctx, ownerSessionID, limit) + if err != nil { + return nil, oops.In("workflow").Code("list_runs").Wrapf(err, "list workflow runs") + } + + return runs, nil +} + +// Events returns durable workflow events in replay order. +func (service *Service) Events( + ctx context.Context, + runID string, + after int64, + limit int, +) ([]database.TaskEventEntity, error) { + events, err := service.runs.Tasks().ListEvents(ctx, runID, after, limit) + if err != nil { + return nil, oops.In("workflow").Code("list_events").Wrapf(err, "list workflow events") + } + + return events, nil +} + +// AgentTask returns one agent task linked from a workflow. +func (service *Service) AgentTask(ctx context.Context, taskID string) (*database.AgentTaskEntity, bool, error) { + task, found, err := service.runs.AgentTasks().Get(ctx, taskID) + if err != nil { + return nil, false, oops.In("workflow").Code("get_agent_task").Wrapf(err, "get workflow agent task") + } + + return task, found, nil +} + +// AgentTasks returns workflow child links in launch order. +func (service *Service) AgentTasks( + ctx context.Context, + runID string, +) ([]database.WorkflowAgentTaskEntity, error) { + tasks, err := service.runs.ListAgentTasks(ctx, runID) + if err != nil { + return nil, oops.In("workflow").Code("list_agent_tasks").Wrapf(err, "list workflow agent tasks") + } + + return tasks, nil +} + +func (service *Service) createRun( + ctx context.Context, + request *ServiceRequest, +) (*database.WorkflowRunEntity, error) { + if request == nil { + return nil, oops.In("workflow").Code("invalid_service_request").Errorf("request is required") + } + + sourceVersion := request.SourceVersion + if sourceVersion == "" { + sourceVersion = defaultSourceVersion + } + + argumentsJSON := request.ArgumentsJSON + if argumentsJSON == "" { + argumentsJSON = "{}" + } + + hash := sha256.Sum256([]byte(request.Source)) + + run, err := service.runs.Create(ctx, &database.WorkflowRunEntity{ + Task: database.TaskEntity{ + CreatedAt: time.Time{}, StartedAt: nil, FinishedAt: nil, UpdatedAt: time.Time{}, + LeaseExpiresAt: nil, ID: "", Kind: "", ParentTaskID: "", OwnerSessionID: request.OwnerSessionID, + ConcurrencyKey: request.OwnerSessionID, LeaseOwner: "", State: "", Result: "", + ErrorCode: "", ErrorMessage: "", + }, + Name: request.Name, Source: request.Source, SourceHash: hex.EncodeToString(hash[:]), + SourceVersion: sourceVersion, ArgumentsJSON: argumentsJSON, + }) + if err != nil { + return nil, oops.In("workflow").Code("create_run").Wrapf(err, "create workflow run") + } + + return run, nil +} + +func decodeArguments(argumentsJSON string) (map[string]any, error) { + arguments := map[string]any{} + if err := json.Unmarshal([]byte(argumentsJSON), &arguments); err != nil { + return nil, oops.In("workflow").Code("decode_arguments").Wrapf(err, "decode workflow arguments") + } + + return arguments, nil +} + +func (service *Service) loadVerifiedRun( + ctx context.Context, + runID string, +) (*database.WorkflowRunEntity, error) { + run, found, err := service.runs.Get(ctx, runID) + if err != nil { + return nil, oops.In("workflow").Code("load_persisted_run").Wrapf(err, "load persisted workflow run") + } + + if !found { + return nil, oops.In("workflow").Code("missing_persisted_run").Errorf("persisted workflow run was not found") + } + + hash := sha256.Sum256([]byte(run.Source)) + if hex.EncodeToString(hash[:]) != run.SourceHash { + return nil, oops.In("workflow").Code("source_hash_mismatch").Errorf("persisted workflow source hash differs") + } + + return run, nil +} + +func (service *Service) eventSink(runID string) EventSink { + return func(ctx context.Context, event Event) error { + payload, err := json.Marshal(eventPayload(event)) + if err != nil { + return oops.In("workflow").Code("encode_event").Wrapf(err, "encode workflow event") + } + + if _, err := service.runs.Tasks().AppendEvent(ctx, runID, workflowEventKind, string(payload)); err != nil { + return oops.In("workflow").Code("persist_event").Wrapf(err, "persist workflow event") + } + + return nil + } +} + +func (service *Service) finish(ctx context.Context, runID string, result *RunResult, runErr error) error { + encoded, encodeErr := json.Marshal(runResultPayload{ + Value: result.Value, Stdout: result.Stdout, Stderr: result.Stderr, + LaunchedTaskIDs: result.LaunchedTaskIDs, TaskResults: result.TaskResults, + }) + if encodeErr != nil { + runErr = errors.Join(runErr, encodeErr) + } + + state, eventKind, errorCode, errorMessage := workflowOutcome(ctx, runErr) + + changed, err := service.runs.Tasks().Finish(context.WithoutCancel(ctx), &database.TaskFinish{ + TaskID: runID, EventKind: eventKind, Result: string(encoded), ErrorCode: errorCode, + ErrorMessage: errorMessage, PayloadJSON: "{}", LeaseOwner: service.leaseOwner, TargetState: state, + From: []database.TaskState{database.TaskRunning, database.TaskCanceling}, + }) + if err != nil { + return oops.In("workflow").Code("finish_run").Wrapf(err, "finish workflow run") + } + + if !changed { + return oops.In("workflow").Code("finish_run_conflict").Errorf("workflow run was not running") + } + + if encodeErr != nil { + return oops.In("workflow").Code("encode_result").Wrapf(encodeErr, "encode workflow result") + } + + return nil +} + +func workflowOutcome( + ctx context.Context, + runErr error, +) (state database.TaskState, eventKind, errorCode, errorMessage string) { + if errors.Is(ctx.Err(), context.Canceled) || errors.Is(runErr, context.Canceled) { + return database.TaskCanceled, workflowCanceled, "canceled", errorString(runErr) + } + + if runErr != nil { + return database.TaskFailed, workflowFailed, "workflow_failed", runErr.Error() + } + + return database.TaskSucceeded, workflowSucceeded, "", "" +} + +func errorString(err error) string { + if err == nil { + return "" + } + + return err.Error() +} + +func (service *Service) monitorLease( + ctx context.Context, + runID string, + cancel context.CancelFunc, +) <-chan error { + result := make(chan error, 1) + + go func() { + defer close(result) + + ticker := time.NewTicker(workflowLeaseInterval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + result <- nil + + return + case <-ticker.C: + stopped, err := service.monitorLeaseTick(ctx, runID) + if stopped { + result <- nil + + return + } + + if err != nil { + cancel() + + result <- err + + return + } + } + } + }() + + return result +} + +func (service *Service) monitorLeaseTick(ctx context.Context, runID string) (stopped bool, err error) { + if contextDone(ctx) { + return true, nil + } + + err = service.renewLease(ctx, runID) + + return contextDone(ctx), err +} + +func contextDone(ctx context.Context) bool { + select { + case <-ctx.Done(): + return true + default: + return false + } +} + +func (service *Service) renewLease(ctx context.Context, runID string) error { + renewed, err := service.runs.Tasks().RenewLease( + ctx, runID, service.leaseOwner, time.Now().UTC().Add(workflowLeaseDuration), + ) + if err != nil { + return oops.In("workflow").Code("renew_lease").Wrapf(err, "renew workflow lease") + } + + if !renewed { + return oops.In("workflow").Code("lease_lost").Errorf("workflow lease was lost") + } + + task, found, err := service.runs.Tasks().Get(ctx, runID) + if err != nil { + return oops.In("workflow").Code("read_lease_state").Wrapf(err, "read workflow state") + } + + if !found { + return oops.In("workflow").Code("missing_lease_task").Errorf("workflow lease task was not found") + } + + if task.State == database.TaskCanceling { + return context.Canceled + } + + return nil +} + +func (service *Service) track(runID string, cancel context.CancelFunc) { + service.mu.Lock() + defer service.mu.Unlock() + + service.active[runID] = cancel +} + +func (service *Service) untrack(runID string) { + service.mu.Lock() + defer service.mu.Unlock() + + delete(service.active, runID) +} diff --git a/internal/workflow/service_internal_test.go b/internal/workflow/service_internal_test.go new file mode 100644 index 00000000..d17d30e8 --- /dev/null +++ b/internal/workflow/service_internal_test.go @@ -0,0 +1,236 @@ +package workflow + +import ( + "context" + "database/sql" + "errors" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + _ "modernc.org/sqlite" // Register the sqlite database/sql driver used by internal service tests. + + "github.com/omarluq/librecode/internal/database" +) + +const serviceStartedEvent = "started" + +type serviceTestController struct{} + +func (serviceTestController) Submit(context.Context, *AgentRequest) (*database.AgentTaskEntity, error) { + return nil, errors.New("unexpected agent submission") +} + +func (serviceTestController) Get(context.Context, string) (*database.AgentTaskEntity, bool, error) { + return nil, false, nil +} + +func (serviceTestController) Await(context.Context, string) (*database.AgentTaskEntity, error) { + return nil, errors.New("unexpected agent await") +} + +func (serviceTestController) Cancel(context.Context, string, string) (*database.TaskEntity, bool, error) { + return nil, false, nil +} + +func TestServiceDefaultsAndHelperBoundaries(t *testing.T) { + t.Parallel() + + service, _, owner, _ := newInternalWorkflowService(t) + run, err := service.Submit(t.Context(), &ServiceRequest{ + Name: "defaults", Source: "1", SourceVersion: "", ArgumentsJSON: "", OwnerSessionID: owner, + }) + require.NoError(t, err) + assert.Equal(t, defaultSourceVersion, run.SourceVersion) + assert.Equal(t, "{}", run.ArgumentsJSON) + + assert.False(t, terminal(database.TaskState("unknown"))) + assert.Empty(t, errorString(nil)) +} + +func TestServiceRenewLeaseStateBoundaries(t *testing.T) { + t.Parallel() + + service, repository, owner, _ := newInternalWorkflowService(t) + + queued, err := service.Submit(t.Context(), serviceTestRequest("queued", owner)) + require.NoError(t, err) + err = service.renewLease(t.Context(), queued.Task.ID) + require.ErrorContains(t, err, "lease was lost") + + running, err := service.Submit(t.Context(), serviceTestRequest("running", owner)) + require.NoError(t, err) + claimed, err := repository.Tasks().ClaimQueued(t.Context(), &database.TaskClaim{ + TaskID: running.Task.ID, LeaseOwner: service.leaseOwner, EventKind: serviceStartedEvent, + LeaseExpiresAt: time.Now().Add(time.Minute), + }) + require.NoError(t, err) + require.True(t, claimed) + require.NoError(t, service.renewLease(t.Context(), running.Task.ID)) + + changed, err := repository.Tasks().Transition(t.Context(), running.Task.ID, + []database.TaskState{database.TaskRunning}, database.TaskCanceling, "canceling") + require.NoError(t, err) + require.True(t, changed) + err = service.renewLease(t.Context(), running.Task.ID) + assert.ErrorIs(t, err, context.Canceled) +} + +func TestServiceRejectsMissingAndMalformedPersistedRuns(t *testing.T) { + t.Parallel() + + service, _, _, _ := newInternalWorkflowService(t) + + _, err := service.ExecuteQueued(t.Context(), "missing") + require.ErrorContains(t, err, "was not found") + _, _, err = service.Resume(t.Context(), "missing") + require.ErrorContains(t, err, "was not found") + + for _, test := range []struct { + name string + resume bool + }{ + {name: "execute queued", resume: false}, + {name: "resume interrupted", resume: true}, + } { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + service, repository, owner, connection := newInternalWorkflowService(t) + run, createErr := service.Submit(t.Context(), serviceTestRequest(test.name, owner)) + require.NoError(t, createErr) + _, updateErr := connection.ExecContext(t.Context(), + `UPDATE workflow_runs SET arguments_json = ? WHERE task_id = ?`, `{"broken":`, run.Task.ID) + require.NoError(t, updateErr) + + if test.resume { + changed, transitionErr := repository.Tasks().Transition(t.Context(), run.Task.ID, + []database.TaskState{database.TaskQueued}, database.TaskInterrupted, "interrupted") + require.NoError(t, transitionErr) + require.True(t, changed) + _, _, createErr = service.Resume(t.Context(), run.Task.ID) + } else { + _, createErr = service.ExecuteQueued(t.Context(), run.Task.ID) + } + + require.ErrorContains(t, createErr, "decode workflow arguments") + }) + } +} + +func TestServiceRenewLeaseDetectsDeletedTask(t *testing.T) { + t.Parallel() + + service, repository, owner, connection := newInternalWorkflowService(t) + run, err := service.Submit(t.Context(), serviceTestRequest("leased", owner)) + require.NoError(t, err) + claimed, err := repository.Tasks().ClaimQueued(t.Context(), &database.TaskClaim{ + TaskID: run.Task.ID, LeaseOwner: service.leaseOwner, EventKind: serviceStartedEvent, + LeaseExpiresAt: time.Now().Add(time.Minute), + }) + require.NoError(t, err) + require.True(t, claimed) + _, err = connection.ExecContext(t.Context(), ` + CREATE TRIGGER delete_renewed_task AFTER UPDATE OF lease_expires_at ON tasks + BEGIN DELETE FROM tasks WHERE id = NEW.id; END`) + require.NoError(t, err) + + err = service.renewLease(t.Context(), run.Task.ID) + require.ErrorContains(t, err, "lease task was not found") +} + +func TestServiceFinishEncodingAndStateConflicts(t *testing.T) { + t.Parallel() + + service, repository, owner, _ := newInternalWorkflowService(t) + queued, err := service.Submit(t.Context(), serviceTestRequest("queued", owner)) + require.NoError(t, err) + err = service.finish(t.Context(), queued.Task.ID, serviceTestResult(nil), nil) + require.ErrorContains(t, err, "was not running") + + running, err := service.Submit(t.Context(), serviceTestRequest("encoding", owner)) + require.NoError(t, err) + claimed, err := repository.Tasks().ClaimQueued(t.Context(), &database.TaskClaim{ + TaskID: running.Task.ID, LeaseOwner: service.leaseOwner, EventKind: serviceStartedEvent, + LeaseExpiresAt: time.Now().Add(time.Minute), + }) + require.NoError(t, err) + require.True(t, claimed) + + err = service.finish(t.Context(), running.Task.ID, serviceTestResult(make(chan int)), nil) + require.ErrorContains(t, err, "encode workflow result") + stored, found, getErr := repository.Get(t.Context(), running.Task.ID) + require.NoError(t, getErr) + require.True(t, found) + assert.Equal(t, database.TaskFailed, stored.Task.State) +} + +func TestServiceRepositoryFailurePaths(t *testing.T) { + t.Parallel() + + service, _, owner, connection := newInternalWorkflowService(t) + run, err := service.Submit(t.Context(), serviceTestRequest("closed", owner)) + require.NoError(t, err) + require.NoError(t, connection.Close()) + + _, _, err = service.Get(t.Context(), run.Task.ID) + require.ErrorContains(t, err, "get workflow run") + _, err = service.List(t.Context(), owner, 10) + require.ErrorContains(t, err, "list workflow runs") + _, err = service.Events(t.Context(), run.Task.ID, 0, 10) + require.ErrorContains(t, err, "list workflow events") + _, _, err = service.AgentTask(t.Context(), run.Task.ID) + require.ErrorContains(t, err, "get workflow agent task") + _, err = service.AgentTasks(t.Context(), run.Task.ID) + require.ErrorContains(t, err, "list workflow agent tasks") + _, err = service.RecoverInterrupted(t.Context()) + require.ErrorContains(t, err, "recover interrupted workflow runs") + _, err = service.Cancel(t.Context(), owner, run.Task.ID) + require.ErrorContains(t, err, "get workflow run") + _, err = service.loadVerifiedRun(t.Context(), run.Task.ID) + require.ErrorContains(t, err, "load persisted workflow run") + err = service.eventSink(run.Task.ID)(t.Context(), Event{ + Task: TaskResult{ID: "", State: "", Result: "", ErrorCode: "", ErrorMessage: ""}, + Kind: EventTaskLaunched, TaskID: "", NodeKey: "", InvocationIndex: 0, + }) + require.ErrorContains(t, err, "persist workflow event") + err = service.renewLease(t.Context(), run.Task.ID) + require.ErrorContains(t, err, "renew workflow lease") +} + +func serviceTestRequest(name, owner string) *ServiceRequest { + return &ServiceRequest{ + Name: name, Source: "1", SourceVersion: "", ArgumentsJSON: "{}", OwnerSessionID: owner, + } +} + +func serviceTestResult(value any) *RunResult { + return &RunResult{ + Value: value, Stdout: "", Stderr: "", LaunchedTaskIDs: nil, TaskResults: nil, + } +} + +func newInternalWorkflowService( + t *testing.T, +) (*Service, *database.WorkflowRepository, string, *sql.DB) { + t.Helper() + + connection, err := sql.Open("sqlite", "file:"+t.Name()+"?mode=memory&cache=shared&_pragma=foreign_keys(1)") + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, connection.Close()) }) + require.NoError(t, database.Migrate(t.Context(), connection)) + + session, err := database.NewSessionRepository(connection).CreateSession( + t.Context(), t.TempDir(), "workflow owner", "", + ) + require.NoError(t, err) + runner, err := NewRunner(serviceTestController{}) + require.NoError(t, err) + + repository := database.NewWorkflowRepository(connection) + service, err := NewService(repository, runner) + require.NoError(t, err) + + return service, repository, session.ID, connection +} diff --git a/internal/workflow/service_test.go b/internal/workflow/service_test.go new file mode 100644 index 00000000..e48cf028 --- /dev/null +++ b/internal/workflow/service_test.go @@ -0,0 +1,524 @@ +package workflow_test + +import ( + "context" + "crypto/sha256" + "database/sql" + "encoding/hex" + "errors" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + _ "modernc.org/sqlite" // Register the sqlite database/sql driver used by service tests. + + "github.com/omarluq/librecode/internal/agent" + "github.com/omarluq/librecode/internal/agenttask" + "github.com/omarluq/librecode/internal/assistant" + "github.com/omarluq/librecode/internal/database" + "github.com/omarluq/librecode/internal/workflow" +) + +const serviceQueuedName = "queued" + +const serviceTestSource = "1 + 1" + +func TestServiceRejectsMissingDependencies(t *testing.T) { + t.Parallel() + + runner, err := workflow.NewRunner(newFakeController()) + require.NoError(t, err) + _, repository, _ := newWorkflowService(t, newFakeController()) + + tests := []struct { + name string + missingRepository bool + }{ + {name: "repository", missingRepository: true}, + {name: "runner", missingRepository: false}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + candidateRepository, candidateRunner := repository, runner + if test.missingRepository { + candidateRepository = nil + } else { + candidateRunner = nil + } + + service, serviceErr := workflow.NewService(candidateRepository, candidateRunner) + require.ErrorContains(t, serviceErr, "repository and runner are required") + assert.Nil(t, service) + }) + } +} + +func TestServiceSubmitRejectsNilRequest(t *testing.T) { + t.Parallel() + + service, _, _ := newWorkflowService(t, newFakeController()) + run, err := service.Submit(t.Context(), nil) + require.ErrorContains(t, err, "request is required") + assert.Nil(t, run) +} + +func TestServiceAwaitMissingRun(t *testing.T) { + t.Parallel() + + service, _, _ := newWorkflowService(t, newFakeController()) + completed, err := service.Await(t.Context(), testUUID()) + require.ErrorContains(t, err, "was not found") + assert.Nil(t, completed) +} + +func TestServiceCancelThenAwait(t *testing.T) { + t.Parallel() + + service, _, owner := newWorkflowService(t, newFakeController()) + run := submitQueuedWorkflow(t, service, owner) + + canceled, err := service.Cancel(t.Context(), owner, run.Task.ID) + require.NoError(t, err) + require.True(t, canceled) + completed, err := service.Await(t.Context(), run.Task.ID) + require.NoError(t, err) + require.NotNil(t, completed) + assert.Equal(t, database.TaskCanceled, completed.Task.State) +} + +func TestServiceCancelMissingRun(t *testing.T) { + t.Parallel() + + service, _, owner := newWorkflowService(t, newFakeController()) + canceled, err := service.Cancel(t.Context(), owner, testUUID()) + require.NoError(t, err) + assert.False(t, canceled) +} + +func TestServiceAccessorsReturnQueuedRunWithoutAgentTasks(t *testing.T) { + t.Parallel() + + service, _, owner := newWorkflowService(t, newFakeController()) + run := submitQueuedWorkflow(t, service, owner) + + runs, err := service.List(t.Context(), owner, 0) + require.NoError(t, err) + assert.Len(t, runs, 1) + links, err := service.AgentTasks(t.Context(), run.Task.ID) + require.NoError(t, err) + assert.Empty(t, links) +} + +func TestServiceAgentTaskReturnsMissing(t *testing.T) { + t.Parallel() + + service, _, _ := newWorkflowService(t, newFakeController()) + task, found, err := service.AgentTask(t.Context(), testUUID()) + require.NoError(t, err) + assert.False(t, found) + assert.Nil(t, task) +} + +func TestServiceRecoverInterruptedWhenNoneExist(t *testing.T) { + t.Parallel() + + service, _, _ := newWorkflowService(t, newFakeController()) + recovered, err := service.RecoverInterrupted(t.Context()) + require.NoError(t, err) + assert.Empty(t, recovered) +} + +func TestServiceAwaitHonorsCancellation(t *testing.T) { + t.Parallel() + + service, _, owner := newWorkflowService(t, newFakeController()) + run, err := service.Submit(t.Context(), &workflow.ServiceRequest{ + Name: serviceQueuedName, Source: serviceTestSource, ArgumentsJSON: "{}", OwnerSessionID: owner, + SourceVersion: "", + }) + require.NoError(t, err) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + _, err = service.Await(ctx, run.Task.ID) + require.Error(t, err) + assert.ErrorIs(t, err, context.Canceled) +} + +func TestServiceRejectsInvalidArguments(t *testing.T) { + t.Parallel() + + service, _, owner := newWorkflowService(t, newFakeController()) + run, result, err := service.Run(t.Context(), &workflow.ServiceRequest{ + Name: "bad arguments", Source: serviceTestSource, ArgumentsJSON: `{"x":`, OwnerSessionID: owner, + SourceVersion: "", + }) + require.Error(t, err) + assert.Nil(t, run) + assert.Nil(t, result) +} + +func TestServiceExecuteQueuedClaimsOnlyOnce(t *testing.T) { + t.Parallel() + + service, _, owner := newWorkflowService(t, newFakeController()) + queued := submitQueuedWorkflow(t, service, owner) + + claimed, err := service.ExecuteQueued(t.Context(), queued.Task.ID) + require.NoError(t, err) + require.True(t, claimed) + claimed, err = service.ExecuteQueued(t.Context(), queued.Task.ID) + require.NoError(t, err) + assert.False(t, claimed) +} + +func testUUID() string { return "018f1234-5678-7abc-8def-0123456789ab" } + +func submitQueuedWorkflow(t *testing.T, service *workflow.Service, owner string) *database.WorkflowRunEntity { + t.Helper() + + run, err := service.Submit(t.Context(), &workflow.ServiceRequest{ + Name: serviceQueuedName, Source: serviceTestSource, ArgumentsJSON: "{}", OwnerSessionID: owner, + SourceVersion: "", + }) + require.NoError(t, err) + + return run +} + +func TestServicePersistsSuccessfulRun(t *testing.T) { + t.Parallel() + + service, repository, owner := newWorkflowService(t, newFakeController()) + run, result, err := service.Run(t.Context(), &workflow.ServiceRequest{ + Name: "inspect", Source: serviceTestSource, SourceVersion: "v1", ArgumentsJSON: "{}", + OwnerSessionID: owner, + }) + require.NoError(t, err) + require.NotNil(t, result) + assert.Equal(t, database.TaskSucceeded, run.Task.State) + assert.Empty(t, result.LaunchedTaskIDs) + + events, err := service.Events(t.Context(), run.Task.ID, 0, 20) + require.NoError(t, err) + assert.Len(t, events, 3) + + stored, found, err := repository.Get(t.Context(), run.Task.ID) + require.NoError(t, err) + require.True(t, found) + assert.Equal(t, database.TaskSucceeded, stored.Task.State) + assert.NotEmpty(t, stored.Task.Result) +} + +func TestServicePersistsFailedRun(t *testing.T) { + t.Parallel() + + service, _, owner := newWorkflowService(t, newFakeController()) + run, _, err := service.Run(t.Context(), &workflow.ServiceRequest{ + Name: "invalid", Source: `func {`, SourceVersion: "v1", ArgumentsJSON: "{}", + OwnerSessionID: owner, + }) + require.Error(t, err) + + stored, found, getErr := service.Get(t.Context(), run.Task.ID) + require.NoError(t, getErr) + require.True(t, found) + assert.Equal(t, database.TaskFailed, stored.Task.State) + assert.Equal(t, "workflow_failed", stored.Task.ErrorCode) +} + +func TestServiceExecuteQueuedPersistsEvaluationFailureWithoutReturningIt(t *testing.T) { + t.Parallel() + + service, _, owner := newWorkflowService(t, newFakeController()) + run, err := service.Submit(t.Context(), &workflow.ServiceRequest{ + Name: "invalid", Source: `func {`, SourceVersion: "v1", ArgumentsJSON: "{}", + OwnerSessionID: owner, + }) + require.NoError(t, err) + + claimed, err := service.ExecuteQueued(t.Context(), run.Task.ID) + require.NoError(t, err) + assert.True(t, claimed) + + stored, found, err := service.Get(t.Context(), run.Task.ID) + require.NoError(t, err) + require.True(t, found) + assert.Equal(t, database.TaskFailed, stored.Task.State) + assert.Equal(t, "workflow_failed", stored.Task.ErrorCode) + assert.Contains(t, stored.Task.ErrorMessage, "evaluate workflow source") +} + +func TestServiceResumesInterruptedRun(t *testing.T) { + t.Parallel() + + service, repository, owner := newWorkflowService(t, newFakeController()) + run, err := repository.Create(t.Context(), verifiedWorkflowRunEntity(owner, "1 + 1")) + require.NoError(t, err) + changed, err := repository.Tasks().Transition(t.Context(), run.Task.ID, + []database.TaskState{database.TaskQueued}, database.TaskInterrupted, "workflow_interrupted") + require.NoError(t, err) + require.True(t, changed) + + resumed, result, err := service.Resume(t.Context(), run.Task.ID) + require.NoError(t, err) + assert.Equal(t, database.TaskSucceeded, resumed.Task.State) + assert.Equal(t, 2, result.Value) + + events, err := service.Events(t.Context(), run.Task.ID, 0, 10) + require.NoError(t, err) + assert.Equal(t, "workflow_resumed", events[2].Event.Kind) +} + +func TestServiceResumeVerifiesPersistedSourceHash(t *testing.T) { + t.Parallel() + + service, repository, owner := newWorkflowService(t, newFakeController()) + run, err := repository.Create(t.Context(), workflowRunEntity(owner)) + require.NoError(t, err) + + _, _, err = service.Resume(t.Context(), run.Task.ID) + require.ErrorContains(t, err, "source hash differs") +} + +func TestServiceRejectsCancelFromAnotherOwner(t *testing.T) { + t.Parallel() + + service, repository, owner := newWorkflowService(t, newFakeController()) + run, err := repository.Create(t.Context(), workflowRunEntity(owner)) + require.NoError(t, err) + + canceled, err := service.Cancel(t.Context(), "other-owner", run.Task.ID) + require.NoError(t, err) + assert.False(t, canceled) +} + +func TestServiceIntegratesDurableAgentTaskLifecycle(t *testing.T) { + t.Parallel() + + environment := newWorkflowIntegration(t) + outcome, taskID, runID := startIntegrationWorkflow(t, environment, "durable-agent") + + task, found, err := environment.agentTasks.Get(t.Context(), taskID) + require.NoError(t, err) + require.True(t, found) + assert.Equal(t, database.TaskRunning, task.Task.State) + assert.Equal(t, runID, task.Task.ParentTaskID) + + require.Eventually(t, func() bool { + links, linkErr := environment.workflows.AgentTasks(t.Context(), runID) + + return linkErr == nil && len(links) == 1 && links[0].AgentTaskID == taskID + }, time.Second, 10*time.Millisecond) + + environment.runner.unblock() + + completed := <-outcome + require.NoError(t, completed.err) + require.NotNil(t, completed.run) + require.NotNil(t, completed.result) + assert.Equal(t, database.TaskSucceeded, completed.run.Task.State) + assert.Equal(t, []string{taskID}, completed.result.LaunchedTaskIDs) + require.Len(t, completed.result.TaskResults, 1) + assert.Equal(t, "integrated result", completed.result.TaskResults[0].Result) + + persisted, found, err := environment.agentTasks.Get(t.Context(), taskID) + require.NoError(t, err) + require.True(t, found) + assert.Equal(t, database.TaskSucceeded, persisted.Task.State) + assert.Equal(t, "integrated result", persisted.Task.Result) +} + +func TestServiceCancellationCascadesToDurableAgentTask(t *testing.T) { + t.Parallel() + + environment := newWorkflowIntegration(t) + outcome, taskID, runID := startIntegrationWorkflow(t, environment, "cancel-agent") + + canceled, err := environment.workflows.Cancel(t.Context(), environment.owner, runID) + require.NoError(t, err) + require.True(t, canceled) + + completed := <-outcome + require.Error(t, completed.err) + require.NotNil(t, completed.run) + assert.Equal(t, database.TaskCanceled, completed.run.Task.State) + + agentResult, err := environment.tasks.Await(t.Context(), taskID) + require.NoError(t, err) + assert.Equal(t, database.TaskCanceled, agentResult.Task.State) +} + +func newWorkflowService( + t *testing.T, + controller workflow.Controller, +) (*workflow.Service, *database.WorkflowRepository, string) { + t.Helper() + + connection, err := sql.Open("sqlite", "file:"+t.Name()+"?mode=memory&cache=shared&_pragma=foreign_keys(1)") + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, connection.Close()) }) + require.NoError(t, database.Migrate(t.Context(), connection)) + + repository := database.NewWorkflowRepository(connection) + sessions := database.NewSessionRepository(connection) + owner, err := sessions.CreateSession(t.Context(), "/work", "workflow owner", "") + require.NoError(t, err) + + runner, err := workflow.NewRunner(controller) + require.NoError(t, err) + service, err := workflow.NewService(repository, runner) + require.NoError(t, err) + + return service, repository, owner.ID +} + +func verifiedWorkflowRunEntity(owner, source string) *database.WorkflowRunEntity { + run := workflowRunEntity(owner) + hash := sha256.Sum256([]byte(source)) + run.Source = source + run.SourceHash = hex.EncodeToString(hash[:]) + + return run +} + +func workflowRunEntity(owner string) *database.WorkflowRunEntity { + return &database.WorkflowRunEntity{ + Task: database.TaskEntity{ + CreatedAt: databaseZeroTime(), StartedAt: nil, FinishedAt: nil, UpdatedAt: databaseZeroTime(), + LeaseExpiresAt: nil, ID: "", Kind: "", ParentTaskID: "", OwnerSessionID: owner, + ConcurrencyKey: owner, LeaseOwner: "", State: "", Result: "", ErrorCode: "", ErrorMessage: "", + }, + Name: "test workflow", Source: "1", SourceHash: "hash", SourceVersion: "v1", ArgumentsJSON: "{}", + } +} + +func databaseZeroTime() (zeroTime time.Time) { + return zeroTime +} + +type integrationAgentRunner struct { + started chan string + release chan struct{} + once sync.Once +} + +func (runner *integrationAgentRunner) Run( + ctx context.Context, + task *database.AgentTaskEntity, + _ agenttask.EventSink, +) (agenttask.Result, error) { + runner.started <- task.Task.ID + + select { + case <-ctx.Done(): + return agenttask.Result{}, errors.Join(errors.New("integration agent canceled"), ctx.Err()) + case <-runner.release: + return agenttask.Result{Text: "integrated result", UsageJSON: `{}`}, nil + } +} + +func (runner *integrationAgentRunner) unblock() { + runner.once.Do(func() { close(runner.release) }) +} + +type workflowIntegration struct { + workflows *workflow.Service + tasks *agenttask.Service + agentTasks *database.AgentTaskRepository + runner *integrationAgentRunner + owner string +} + +type workflowRunOutcome struct { + run *database.WorkflowRunEntity + result *workflow.RunResult + err error +} + +func newWorkflowIntegration(t *testing.T) *workflowIntegration { + t.Helper() + + connection, err := sql.Open("sqlite", "file:"+t.Name()+"?mode=memory&cache=shared&_pragma=foreign_keys(1)") + require.NoError(t, err) + connection.SetMaxOpenConns(1) + t.Cleanup(func() { require.NoError(t, connection.Close()) }) + require.NoError(t, database.Migrate(t.Context(), connection)) + + taskRepository := database.NewTaskRepository(connection) + agentTaskRepository := database.NewAgentTaskRepository(connection) + sessions := database.NewSessionRepository(connection) + workflowRepository := database.NewWorkflowRepository(connection) + owner, err := sessions.CreateSession(t.Context(), t.TempDir(), "workflow integration owner", "") + require.NoError(t, err) + + agentRunner := &integrationAgentRunner{ + started: make(chan string, 1), release: make(chan struct{}), once: sync.Once{}, + } + taskService, err := agenttask.New(context.Background(), &agenttask.Options{ + Tasks: taskRepository, AgentTasks: agentTaskRepository, Workflows: workflowRepository, + Runner: agentRunner, Logger: nil, + Concurrency: 1, SessionConcurrency: 1, QueueCapacity: 4, Timeout: time.Minute, + }) + require.NoError(t, err) + t.Cleanup(func() { + agentRunner.unblock() + require.NoError(t, taskService.Shutdown(context.Background())) + }) + + submitter, err := assistant.NewAgentSubmitter(taskService, agent.Load(t.TempDir())) + require.NoError(t, err) + controller, err := assistant.NewWorkflowController(submitter, taskService, sessions) + require.NoError(t, err) + workflowRunner, err := workflow.NewRunner(controller) + require.NoError(t, err) + workflowService, err := workflow.NewService(workflowRepository, workflowRunner) + require.NoError(t, err) + + return &workflowIntegration{ + workflows: workflowService, tasks: taskService, agentTasks: agentTaskRepository, + runner: agentRunner, owner: owner.ID, + } +} + +func startIntegrationWorkflow( + t *testing.T, + environment *workflowIntegration, + name string, +) (outcome <-chan workflowRunOutcome, taskID, runID string) { + t.Helper() + + outcomeChannel := make(chan workflowRunOutcome, 1) + + go func() { + run, result, err := environment.workflows.Run(context.Background(), &workflow.ServiceRequest{ + Name: name, Source: agentSource, SourceVersion: "", ArgumentsJSON: "{}", + OwnerSessionID: environment.owner, + }) + outcomeChannel <- workflowRunOutcome{run: run, result: result, err: err} + }() + + taskID = awaitIntegrationTask(t, environment.runner.started) + runs, err := environment.workflows.List(t.Context(), environment.owner, 10) + require.NoError(t, err) + require.Len(t, runs, 1) + + return outcomeChannel, taskID, runs[0].Task.ID +} + +func awaitIntegrationTask(t *testing.T, started <-chan string) string { + t.Helper() + + select { + case taskID := <-started: + return taskID + case <-time.After(10 * time.Second): + require.FailNow(t, "timed out waiting for integrated agent task to start") + + return "" + } +} diff --git a/internal/workflow/worker_testmain_test.go b/internal/workflow/worker_testmain_test.go new file mode 100644 index 00000000..8164bd0a --- /dev/null +++ b/internal/workflow/worker_testmain_test.go @@ -0,0 +1,20 @@ +package workflow_test + +import ( + "os" + "testing" + + "github.com/omarluq/librecode/internal/executeworker" +) + +func TestMain(testMain *testing.M) { + if len(os.Args) == 2 && os.Args[1] == "__execute-worker" { + if err := executeworker.Serve(os.Stdin, os.Stdout); err != nil { + os.Exit(2) + } + + os.Exit(0) + } + + os.Exit(testMain.Run()) +} diff --git a/internal/workflow/workflow.go b/internal/workflow/workflow.go new file mode 100644 index 00000000..8a818d61 --- /dev/null +++ b/internal/workflow/workflow.go @@ -0,0 +1,620 @@ +// Package workflow executes MVM source that coordinates durable agent tasks. +package workflow + +import ( + "context" + "encoding/json" + "errors" + "strings" + "sync" + "time" + + "github.com/samber/oops" + + "github.com/omarluq/librecode/internal/database" + "github.com/omarluq/librecode/internal/executeworker" +) + +const ( + cancelTimeout = 5 * time.Second + pipelineResultKind = "pipeline_result" + taskNotOwnedBySession = "task is not owned by this workflow session" +) + +// EventKind identifies workflow progress. +type EventKind string + +const ( + // EventTaskLaunched reports a newly submitted task. + EventTaskLaunched EventKind = "task_launched" + // EventTaskCompleted reports a completed wait. + EventTaskCompleted EventKind = "task_completed" +) + +// Event is an observable workflow progress update. +type Event struct { + Task TaskResult + Kind EventKind + TaskID string + NodeKey string + InvocationIndex int +} + +// EventSink receives workflow progress synchronously. +type EventSink func(context.Context, Event) error + +// AgentOptions configures one durable agent launch. The controller supplies any +// product-specific defaults not provided by the workflow. +type AgentOptions struct { + NodeKey string `json:"node_key"` + AgentName string `json:"agent_name"` + Model string `json:"model"` + Provider string `json:"provider"` + ConcurrencyKey string `json:"concurrency_key"` + Depth int `json:"depth"` +} + +// AgentRequest describes a durable agent operation requested by a workflow. +type AgentRequest struct { + ParentTaskID string + OwnerSessionID string + NodeKey string + Prompt string + Options AgentOptions + InvocationIndex int +} + +// Controller is the narrow durable-agent boundary used by workflow runs. +type Controller interface { + Submit(context.Context, *AgentRequest) (*database.AgentTaskEntity, error) + Get(context.Context, string) (*database.AgentTaskEntity, bool, error) + Await(context.Context, string) (*database.AgentTaskEntity, error) + Cancel(context.Context, string, string) (*database.TaskEntity, bool, error) +} + +// RunRequest describes one isolated workflow evaluation. +type RunRequest struct { + RunID string + Name string + Source string + OwnerSessionID string + OnEvent EventSink + Arguments map[string]any + PersistedLinks []database.WorkflowAgentTaskEntity +} + +// RunResult contains script output and the tasks launched by this run. +type RunResult struct { + Value any + Stdout string + Stderr string + LaunchedTaskIDs []string + TaskResults []TaskResult +} + +// TaskResult is the stable workflow-facing view of a durable task. +type TaskResult struct { + ID string `json:"id"` + State string `json:"state"` + Result string `json:"result"` + ErrorCode string `json:"error_code"` + ErrorMessage string `json:"error_message"` +} + +// Runner evaluates workflow source against a durable agent controller. +type Runner struct { + controller Controller + executable string +} + +// NewRunner creates a workflow runner. +func NewRunner(controller Controller) (*Runner, error) { + if controller == nil { + return nil, oops.In("workflow").Code("invalid_controller").Errorf("controller is required") + } + + return &Runner{controller: controller, executable: ""}, nil +} + +// Run evaluates source and waits for any Agent/Wait calls made by that source. +func (runner *Runner) Run(ctx context.Context, request *RunRequest) (runResult RunResult, runErr error) { + if ctx == nil { + return RunResult{}, oops.In("workflow").Code("invalid_context").Errorf("context is required") + } + + if request == nil { + return RunResult{}, oops.In("workflow").Code("invalid_request").Errorf("request is required") + } + + run := &runHost{ + runID: request.RunID, ownerSessionID: request.OwnerSessionID, controller: runner.controller, + onEvent: request.OnEvent, launched: make(map[string]struct{}), taskIDs: make([]string, 0), + invocations: make(map[string]int), persisted: make(map[invocationKey]persistedInvocation), + mu: sync.Mutex{}, launchMu: sync.Mutex{}, eventMu: sync.Mutex{}, + } + for _, link := range request.PersistedLinks { + key := invocationKey{nodeKey: normalizeNodeKey(link.NodeKey), index: link.InvocationIndex} + + run.persisted[key] = persistedInvocation{taskID: link.AgentTaskID} + + run.launched[link.AgentTaskID] = struct{}{} + run.taskIDs = append(run.taskIDs, link.AgentTaskID) + } + + client := executeworker.Client{Executable: runner.executable, Handler: run.handleRPC} + result, err := client.EvalRequest(ctx, &executeworker.Request{ + Mode: "workflow", Name: request.Name, Source: request.Source, Arguments: request.Arguments, + }) + + value := normalizeWorkflowValue(result.Value) + if result.ValueKind == pipelineResultKind { + value = normalizePipelineResults(value) + } + + if err != nil { + runErr = errors.Join(runErr, oops.In("workflow").Code("evaluate_source").Wrapf(err, "evaluate workflow source")) + runErr = errors.Join(runErr, run.cancelActive(ctx)) + } + + taskResults, snapshotErr := run.taskResults(ctx) + runResult = RunResult{ + Value: value, Stdout: result.Stdout, Stderr: result.Stderr, + LaunchedTaskIDs: run.launchedTaskIDs(), TaskResults: taskResults, + } + + if snapshotErr != nil { + runErr = errors.Join(runErr, oops.In("workflow").Code("snapshot_tasks"). + Wrapf(snapshotErr, "snapshot workflow tasks")) + if err == nil { + runErr = errors.Join(runErr, run.cancelActive(ctx)) + } + } + + if runErr != nil { + return runResult, runErr + } + + return runResult, nil +} + +func normalizeWorkflowValue(value any) any { + raw, err := json.Marshal(value) + if err != nil { + return value + } + + decoder := json.NewDecoder(strings.NewReader(string(raw))) + decoder.UseNumber() + + var normalized any + if err := decoder.Decode(&normalized); err != nil { + return value + } + + return normalizeNumbers(normalized) +} + +func normalizePipelineResults(value any) any { + items, ok := value.([]any) + if !ok { + return value + } + + if len(items) == 0 { + return []PipelineResult{} + } + + results := make([]PipelineResult, len(items)) + for index, item := range items { + fields, fieldsOK := item.(map[string]any) + if !fieldsOK || len(fields) != 3 { + return value + } + + result, ok := pipelineResultFromFields(fields) + if !ok { + return value + } + + results[index] = result + if results[index].Value == nil && results[index].Error != pipelineNotScheduled { + results[index].Value = TaskResult{ + ID: "", State: "", Result: "", ErrorCode: "", ErrorMessage: "", + } + } + } + + return results +} + +func pipelineResultFromFields(fields map[string]any) (PipelineResult, bool) { + itemIndex, indexOK := fields["index"].(int) + itemValue, valueOK := fields["value"] + itemError, errorOK := fields["error"].(string) + + return PipelineResult{Index: itemIndex, Value: itemValue, Error: itemError}, indexOK && valueOK && errorOK +} + +func normalizeNumbers(value any) any { + switch typed := value.(type) { + case json.Number: + if integer, err := typed.Int64(); err == nil { + return int(integer) + } + + decimal, err := typed.Float64() + if err != nil { + return typed.String() + } + + return decimal + case []any: + for index := range typed { + typed[index] = normalizeNumbers(typed[index]) + } + case map[string]any: + for key := range typed { + typed[key] = normalizeNumbers(typed[key]) + } + } + + return value +} + +type invocationKey struct { + nodeKey string + index int +} + +type persistedInvocation struct { + taskID string +} + +type workflowAgentRPCInput struct { + Prompt string `json:"prompt"` + Options []AgentOptions `json:"options"` +} + +type runHost struct { + controller Controller + runID string + onEvent EventSink + ownerSessionID string + launched map[string]struct{} + invocations map[string]int + persisted map[invocationKey]persistedInvocation + taskIDs []string + mu sync.Mutex + launchMu sync.Mutex + eventMu sync.Mutex +} + +func (host *runHost) handleRPC(ctx context.Context, message *executeworker.Message) (any, error) { + switch message.Method { + case "workflow_agent": + var input workflowAgentRPCInput + if err := json.Unmarshal(message.Input, &input); err != nil { + return nil, oops.In("workflow").Code("invalid_agent_rpc").Wrapf(err, "decode agent request") + } + + return host.agent(ctx, input.Prompt, input.Options...) + case "workflow_wait": + return host.wait(ctx, message.Name) + case "workflow_list": + return host.list(ctx) + case "workflow_cancel": + return host.cancel(ctx, message.Name) + default: + return nil, oops.In("workflow").Code("invalid_rpc_method").Errorf("unknown workflow RPC %q", message.Method) + } +} + +func (host *runHost) agent(ctx context.Context, prompt string, options ...AgentOptions) (string, error) { + // A pipeline may invoke Agent concurrently. Keep invocation allocation, + // persisted-link replay, submission, and launch publication in one stable + // order so a later invocation cannot be persisted before an earlier one. + host.launchMu.Lock() + defer host.launchMu.Unlock() + + if err := ctx.Err(); err != nil { + return "", oops.In("workflow").Code("run_canceled").Wrapf(err, "launch agent") + } + + agentOptions, nodeKey, err := validatedAgentInput(prompt, options) + if err != nil { + return "", err + } + + host.mu.Lock() + invocationIndex := host.invocations[nodeKey] + host.invocations[nodeKey] = invocationIndex + 1 + host.mu.Unlock() + + if reusedTaskID, reused, reuseErr := host.persistedTask(ctx, nodeKey, invocationIndex); reuseErr != nil || reused { + return reusedTaskID, reuseErr + } + + task, err := host.controller.Submit(ctx, &AgentRequest{ + ParentTaskID: host.runID, + OwnerSessionID: host.ownerSessionID, + NodeKey: nodeKey, + InvocationIndex: invocationIndex, + Prompt: prompt, + Options: agentOptions, + }) + if err != nil { + return "", oops.In("workflow").Code("submit_agent").Wrapf(err, "submit agent task") + } + + if task == nil || task.Task.ID == "" { + return "", oops.In("workflow").Code("invalid_agent_task").Errorf("controller returned an empty task") + } + + if task.Task.OwnerSessionID != host.ownerSessionID { + return "", oops.In("workflow").Code("owner_mismatch"). + Errorf("controller returned a task owned by another session") + } + + host.mu.Lock() + host.launched[task.Task.ID] = struct{}{} + host.taskIDs = append(host.taskIDs, task.Task.ID) + host.mu.Unlock() + + if err := host.emit(ctx, &Event{ + Kind: EventTaskLaunched, TaskID: task.Task.ID, NodeKey: nodeKey, + InvocationIndex: invocationIndex, Task: taskResult(task), + }); err != nil { + return "", err + } + + return task.Task.ID, nil +} + +func (host *runHost) wait(ctx context.Context, taskID string) (TaskResult, error) { + if !host.owns(taskID) { + return TaskResult{}, oops.In("workflow").Code("task_not_owned").Errorf("task was not launched by this workflow") + } + + task, found, err := host.controller.Get(ctx, taskID) + if err != nil { + return TaskResult{}, oops.In("workflow").Code("get_task").Wrapf(err, "get agent task") + } + + if !found || task.Task.OwnerSessionID != host.ownerSessionID { + return TaskResult{}, oops.In("workflow").Code("task_not_owned"). + Errorf(taskNotOwnedBySession) + } + + task, err = host.controller.Await(ctx, taskID) + if err != nil { + return TaskResult{}, oops.In("workflow").Code("await_task").Wrapf(err, "await agent task") + } + + if task == nil || task.Task.OwnerSessionID != host.ownerSessionID { + return TaskResult{}, oops.In("workflow").Code("task_not_owned"). + Errorf("controller returned a task owned by another session") + } + + result := taskResult(task) + if err := host.emit(ctx, &Event{ + Task: result, Kind: EventTaskCompleted, TaskID: taskID, NodeKey: "", InvocationIndex: 0, + }); err != nil { + return TaskResult{}, err + } + + return result, nil +} + +func (host *runHost) list(ctx context.Context) ([]TaskResult, error) { + taskIDs := host.launchedTaskIDs() + + results := make([]TaskResult, 0, len(taskIDs)) + for _, taskID := range taskIDs { + task, found, err := host.controller.Get(ctx, taskID) + if err != nil { + return nil, oops.In("workflow").Code("list_task").Wrapf(err, "get launched task %s", taskID) + } + + if !found || task == nil || task.Task.OwnerSessionID != host.ownerSessionID { + return nil, oops.In("workflow").Code("task_not_owned"). + Errorf(taskNotOwnedBySession) + } + + results = append(results, taskResult(task)) + } + + return results, nil +} + +func (host *runHost) cancel(ctx context.Context, taskID string) (TaskResult, error) { + if !host.owns(taskID) { + return TaskResult{}, oops.In("workflow").Code("task_not_owned"). + Errorf("task was not launched by this workflow") + } + + task, found, err := host.controller.Get(ctx, taskID) + if err != nil { + return TaskResult{}, oops.In("workflow").Code("get_task").Wrapf(err, "get agent task") + } + + if !found || task == nil || task.Task.OwnerSessionID != host.ownerSessionID { + return TaskResult{}, oops.In("workflow").Code("task_not_owned"). + Errorf(taskNotOwnedBySession) + } + + canceled, found, err := host.controller.Cancel(ctx, host.ownerSessionID, taskID) + if err != nil { + return TaskResult{}, oops.In("workflow").Code("cancel_task").Wrapf(err, "cancel agent task") + } + + if !found || canceled == nil || canceled.OwnerSessionID != host.ownerSessionID { + return TaskResult{}, oops.In("workflow").Code("task_not_owned"). + Errorf(taskNotOwnedBySession) + } + + return taskResultFromTask(canceled), nil +} + +func validatedAgentInput(prompt string, options []AgentOptions) (AgentOptions, string, error) { + if prompt == "" { + return AgentOptions{}, "", oops.In("workflow").Code("invalid_agent_prompt"). + Errorf("agent prompt is required") + } + + agentOptions, err := oneAgentOptions(options) + if err != nil { + return AgentOptions{}, "", err + } + + return agentOptions, normalizeNodeKey(agentOptions.NodeKey), nil +} + +func normalizeNodeKey(nodeKey string) string { + nodeKey = strings.TrimSpace(nodeKey) + if nodeKey == "" { + return "agent" + } + + return nodeKey +} + +func oneAgentOptions(options []AgentOptions) (AgentOptions, error) { + switch len(options) { + case 0: + return AgentOptions{ + NodeKey: "", AgentName: "", Model: "", Provider: "", ConcurrencyKey: "", Depth: 0, + }, nil + case 1: + return options[0], nil + default: + return AgentOptions{}, oops.In("workflow").Code("invalid_agent_options"). + Errorf("agent accepts at most one options value") + } +} + +func (host *runHost) persistedTask( + ctx context.Context, + nodeKey string, + invocationIndex int, +) (taskID string, found bool, err error) { + key := invocationKey{nodeKey: normalizeNodeKey(nodeKey), index: invocationIndex} + + host.mu.Lock() + invocation, persisted := host.persisted[key] + host.mu.Unlock() + + if !persisted { + return "", false, nil + } + + task, found, err := host.controller.Get(ctx, invocation.taskID) + if err != nil { + return "", false, oops.In("workflow").Code("get_persisted_task"). + Wrapf(err, "get persisted agent task") + } + + if !found || task == nil { + return "", false, nil + } + + host.mu.Lock() + current, persisted := host.persisted[key] + + if !persisted || current != invocation { + host.mu.Unlock() + + return "", false, nil + } + + delete(host.persisted, key) + host.mu.Unlock() + + return invocation.taskID, true, nil +} + +func (host *runHost) owns(taskID string) bool { + host.mu.Lock() + defer host.mu.Unlock() + + _, found := host.launched[taskID] + + return found +} + +func (host *runHost) launchedTaskIDs() []string { + host.mu.Lock() + defer host.mu.Unlock() + + return append([]string(nil), host.taskIDs...) +} + +func (host *runHost) taskResults(ctx context.Context) ([]TaskResult, error) { + return host.list(context.WithoutCancel(ctx)) +} + +func (host *runHost) emit(ctx context.Context, event *Event) error { + if host.onEvent == nil { + return nil + } + + host.eventMu.Lock() + defer host.eventMu.Unlock() + + if err := host.onEvent(ctx, *event); err != nil { + return oops.In("workflow").Code("emit_event").Wrapf(err, "emit workflow event") + } + + return nil +} + +func (host *runHost) cancelActive(ctx context.Context) error { + ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), cancelTimeout) + defer cancel() + + var cancelErr error + + for _, taskID := range host.launchedTaskIDs() { + task, found, err := host.controller.Get(ctx, taskID) + if err != nil { + wrapped := oops.In("workflow").Code("get_task_for_cancel"). + Wrapf(err, "get launched task %s", taskID) + cancelErr = errors.Join(cancelErr, wrapped) + + continue + } + + if !found || task == nil || terminal(task.Task.State) { + continue + } + + if _, _, err := host.controller.Cancel(ctx, host.ownerSessionID, taskID); err != nil { + wrapped := oops.In("workflow").Code("cancel_task"). + Wrapf(err, "cancel launched task %s", taskID) + cancelErr = errors.Join(cancelErr, wrapped) + } + } + + return cancelErr +} + +func taskResult(task *database.AgentTaskEntity) TaskResult { + return taskResultFromTask(&task.Task) +} + +func taskResultFromTask(task *database.TaskEntity) TaskResult { + return TaskResult{ + ID: task.ID, State: string(task.State), Result: task.Result, + ErrorCode: task.ErrorCode, ErrorMessage: task.ErrorMessage, + } +} + +func terminal(state database.TaskState) bool { + switch state { + case database.TaskSucceeded, database.TaskFailed, database.TaskCanceled, database.TaskInterrupted: + return true + case database.TaskQueued, database.TaskRunning, database.TaskCanceling: + return false + default: + return false + } +} diff --git a/internal/workflow/workflow_internal_test.go b/internal/workflow/workflow_internal_test.go new file mode 100644 index 00000000..19744bc8 --- /dev/null +++ b/internal/workflow/workflow_internal_test.go @@ -0,0 +1,359 @@ +package workflow + +import ( + "context" + "encoding/json" + "errors" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/omarluq/librecode/internal/database" + "github.com/omarluq/librecode/internal/executeworker" +) + +type branchController struct { + submit func(context.Context, *AgentRequest) (*database.AgentTaskEntity, error) + get func(context.Context, string) (*database.AgentTaskEntity, bool, error) + await func(context.Context, string) (*database.AgentTaskEntity, error) + cancel func(context.Context, string, string) (*database.TaskEntity, bool, error) +} + +func (controller *branchController) Submit( + ctx context.Context, + request *AgentRequest, +) (*database.AgentTaskEntity, error) { + return controller.submit(ctx, request) +} + +func (controller *branchController) Get(ctx context.Context, id string) (*database.AgentTaskEntity, bool, error) { + return controller.get(ctx, id) +} + +func (controller *branchController) Await(ctx context.Context, id string) (*database.AgentTaskEntity, error) { + return controller.await(ctx, id) +} + +func (controller *branchController) Cancel(ctx context.Context, owner, id string) (*database.TaskEntity, bool, error) { + return controller.cancel(ctx, owner, id) +} + +func TestRunnerRejectsInvalidConstructionAndRunInputs(t *testing.T) { + t.Parallel() + + runner, err := NewRunner(nil) + require.Error(t, err) + assert.Nil(t, runner) + + runner = &Runner{controller: nil, executable: ""} + + var nilContext context.Context + + _, err = runner.Run(nilContext, &RunRequest{ + RunID: "", Name: "", Source: "", OwnerSessionID: "", OnEvent: nil, + Arguments: nil, PersistedLinks: nil, + }) + require.ErrorContains(t, err, "context is required") + _, err = runner.Run(t.Context(), nil) + require.ErrorContains(t, err, "request is required") +} + +func TestWorkflowValueNormalizationFallbacks(t *testing.T) { + t.Parallel() + + unencodable := make(chan int) + assert.Equal(t, unencodable, normalizeWorkflowValue(unencodable)) + assert.InDelta(t, 1.25, normalizeNumbers(json.Number("1.25")), 0) + assert.Equal(t, "not-a-number", normalizeNumbers(json.Number("not-a-number"))) + + original := map[string]any{"number": json.Number("2"), "items": []any{json.Number("3")}} + assert.Equal(t, map[string]any{"number": 2, "items": []any{3}}, normalizeNumbers(original)) + + assert.Equal(t, "ordinary", normalizePipelineResults("ordinary")) + assert.Equal(t, []PipelineResult{}, normalizePipelineResults([]any{})) + + for _, malformed := range []any{ + []any{"not fields"}, + []any{map[string]any{"index": 0, "value": "x"}}, + []any{map[string]any{"index": "zero", "value": "x", "error": ""}}, + } { + assert.Equal(t, malformed, normalizePipelineResults(malformed)) + } +} + +func TestRunHostRPCValidation(t *testing.T) { + t.Parallel() + + host := newBranchHost(newBranchController()) + message := branchMessage("workflow_agent") + message.Input = []byte("{") + _, err := host.handleRPC(t.Context(), message) + require.ErrorContains(t, err, "decode agent request") + _, err = host.handleRPC(t.Context(), branchMessage("unknown")) + require.ErrorContains(t, err, "unknown workflow RPC") +} + +func TestRunHostAgentFailureBranches(t *testing.T) { + t.Parallel() + + failure := errors.New("controller failure") + tests := []struct { + name string + controller *branchController + want string + }{ + { + name: "submit failure", + controller: &branchController{ + submit: func( + context.Context, *AgentRequest, + ) (*database.AgentTaskEntity, error) { + return nil, failure + }, get: nil, await: nil, cancel: nil}, + want: "submit agent task", + }, + { + name: "empty task", + controller: &branchController{ + submit: func( + context.Context, *AgentRequest, + ) (*database.AgentTaskEntity, error) { + return branchAgentTask("", testBranchOwner, database.TaskRunning), nil + }, get: nil, await: nil, cancel: nil}, + want: "empty task", + }, + { + name: "wrong owner", + controller: &branchController{ + submit: func( + context.Context, *AgentRequest, + ) (*database.AgentTaskEntity, error) { + return branchAgentTask("task", "other", database.TaskRunning), nil + }, get: nil, await: nil, cancel: nil}, + want: "another session", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + host := newBranchHost(test.controller) + _, err := host.agent(t.Context(), "prompt") + require.ErrorContains(t, err, test.want) + }) + } + + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + host := newBranchHost(newBranchController()) + _, err := host.agent(ctx, "prompt") + require.ErrorContains(t, err, "launch agent") + + _, err = newBranchHost(newBranchController()).agent(t.Context(), "") + require.ErrorContains(t, err, "prompt is required") + + options := AgentOptions{ + NodeKey: "", AgentName: "", Model: "", Provider: "", ConcurrencyKey: "", Depth: 0, + } + _, err = newBranchHost(newBranchController()).agent(t.Context(), "prompt", options, options) + require.ErrorContains(t, err, "at most one") +} + +func TestRunHostWaitFailureBranches(t *testing.T) { + t.Parallel() + + failure := errors.New("controller failure") + tests := []struct { + name string + controller *branchController + want string + }{ + { + name: "get failure", + controller: &branchController{ + submit: nil, + get: func( + context.Context, string, + ) (*database.AgentTaskEntity, bool, error) { + return nil, false, failure + }, await: nil, cancel: nil}, + want: "get agent task", + }, + { + name: "not found", + controller: &branchController{ + submit: nil, + get: func( + context.Context, string, + ) (*database.AgentTaskEntity, bool, error) { + return nil, false, nil + }, await: nil, cancel: nil}, + want: taskNotOwnedBySession, + }, + { + name: "await failure", + controller: &branchController{ + submit: nil, get: ownedGet, + await: func(context.Context, string) (*database.AgentTaskEntity, error) { return nil, failure }, + cancel: nil, + }, + want: "await agent task", + }, + { + name: "nil awaited task", + controller: &branchController{ + submit: nil, get: ownedGet, + await: func(context.Context, string) (*database.AgentTaskEntity, error) { + return branchAgentTask("task", "other", database.TaskRunning), nil + }, + cancel: nil, + }, + want: "another session", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + host := newBranchHost(test.controller) + host.launched["task"] = struct{}{} + _, err := host.wait(t.Context(), "task") + require.ErrorContains(t, err, test.want) + }) + } +} + +func TestRunHostListAndCancelFailureBranches(t *testing.T) { + t.Parallel() + + failure := errors.New("controller failure") + controller := newBranchController() + controller.get = func(context.Context, string) (*database.AgentTaskEntity, bool, error) { + return nil, false, failure + } + host := newBranchHost(controller) + host.launched["task"] = struct{}{} + host.taskIDs = []string{"task"} + _, err := host.list(t.Context()) + require.ErrorContains(t, err, "get launched task") + + controller = newBranchController() + controller.get = ownedGet + controller.cancel = func(context.Context, string, string) (*database.TaskEntity, bool, error) { + return nil, false, failure + } + host.controller = controller + _, err = host.cancel(t.Context(), "task") + require.ErrorContains(t, err, "cancel agent task") + + controller = newBranchController() + controller.get = ownedGet + controller.cancel = func(context.Context, string, string) (*database.TaskEntity, bool, error) { + return nil, false, nil + } + host.controller = controller + _, err = host.cancel(t.Context(), "task") + require.ErrorContains(t, err, taskNotOwnedBySession) +} + +func TestRunHostPersistedAndCleanupFailures(t *testing.T) { + t.Parallel() + + failure := errors.New("controller failure") + controller := newBranchController() + controller.get = func(context.Context, string) (*database.AgentTaskEntity, bool, error) { + return nil, false, failure + } + host := newBranchHost(controller) + host.persisted[invocationKey{nodeKey: "agent", index: 0}] = persistedInvocation{taskID: "task"} + _, _, err := host.persistedTask(t.Context(), "agent", 0) + require.ErrorContains(t, err, "get persisted agent task") + + cancelCalls := 0 + host = newBranchHost(&branchController{ + submit: nil, + get: func(_ context.Context, taskID string) (*database.AgentTaskEntity, bool, error) { + if taskID == "get-error" { + return nil, false, failure + } + + switch taskID { + case "terminal": + return branchAgentTask(taskID, testBranchOwner, database.TaskSucceeded), true, nil + default: + return branchAgentTask(taskID, testBranchOwner, database.TaskRunning), true, nil + } + }, + await: nil, + cancel: func(context.Context, string, string) (*database.TaskEntity, bool, error) { + cancelCalls++ + + return nil, false, failure + }, + }) + host.taskIDs = []string{"get-error", "terminal", "running"} + err = host.cancelActive(t.Context()) + require.ErrorContains(t, err, "get launched task get-error") + require.ErrorContains(t, err, "cancel launched task running") + assert.Equal(t, 1, cancelCalls) +} + +func TestTerminalStates(t *testing.T) { + t.Parallel() + + for _, state := range []database.TaskState{ + database.TaskSucceeded, database.TaskFailed, database.TaskCanceled, database.TaskInterrupted, + } { + assert.True(t, terminal(state)) + } + + for _, state := range []database.TaskState{ + database.TaskQueued, database.TaskRunning, database.TaskCanceling, database.TaskState("unknown"), + } { + assert.False(t, terminal(state)) + } +} + +const testBranchOwner = "branch-owner" + +func newBranchController() *branchController { + return &branchController{submit: nil, get: nil, await: nil, cancel: nil} +} + +func branchMessage(method string) *executeworker.Message { + return &executeworker.Message{ + Stderr: "", Source: "", Method: method, Mode: "", Name: "", Query: "", Stdout: "", Type: "", + Error: "", ErrorKind: "", ValueKind: "", Input: nil, Value: nil, Arguments: nil, + ID: 0, ExitCode: 0, + } +} + +func newBranchHost(controller Controller) *runHost { + return &runHost{ + controller: controller, runID: "", onEvent: nil, + ownerSessionID: testBranchOwner, launched: make(map[string]struct{}), + invocations: make(map[string]int), persisted: make(map[invocationKey]persistedInvocation), + taskIDs: make([]string, 0), mu: sync.Mutex{}, launchMu: sync.Mutex{}, eventMu: sync.Mutex{}, + } +} + +func ownedGet(_ context.Context, taskID string) (*database.AgentTaskEntity, bool, error) { + return branchAgentTask(taskID, testBranchOwner, database.TaskRunning), true, nil +} + +func branchAgentTask(id, owner string, state database.TaskState) *database.AgentTaskEntity { + return &database.AgentTaskEntity{ + Task: database.TaskEntity{ + CreatedAt: time.Time{}, StartedAt: nil, FinishedAt: nil, UpdatedAt: time.Time{}, LeaseExpiresAt: nil, + ID: id, Kind: "", ParentTaskID: "", OwnerSessionID: owner, ConcurrencyKey: "", LeaseOwner: "", + State: state, Result: "", ErrorCode: "", ErrorMessage: "", + }, + ChildSessionID: "", AgentName: "", Prompt: "", Model: "", Provider: "", PolicyJSON: "", UsageJSON: "", Depth: 0, + } +} diff --git a/internal/workflow/workflow_test.go b/internal/workflow/workflow_test.go new file mode 100644 index 00000000..0d137f5b --- /dev/null +++ b/internal/workflow/workflow_test.go @@ -0,0 +1,525 @@ +package workflow_test + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/omarluq/librecode/internal/database" + "github.com/omarluq/librecode/internal/workflow" +) + +const ( + testOwner = "session-1" + firstTask = "task-1" + secondTask = "task-2" + agentSource = `import "librecode/workflow"; id, _ := workflow.Agent("inspect"); workflow.Wait(id)` +) + +type fakeController struct { + tasks map[string]*database.AgentTaskEntity + await func(context.Context, string) (*database.AgentTaskEntity, error) + submitCh chan string + taskIDs []string + submits []*workflow.AgentRequest + cancels [][2]string + mu sync.Mutex +} + +func (fake *fakeController) Submit( + _ context.Context, + request *workflow.AgentRequest, +) (*database.AgentTaskEntity, error) { + fake.mu.Lock() + defer fake.mu.Unlock() + + taskID := fmt.Sprintf("task-%d", len(fake.submits)+1) + if len(fake.taskIDs) > len(fake.submits) { + taskID = fake.taskIDs[len(fake.submits)] + } + + copyRequest := *request + fake.submits = append(fake.submits, ©Request) + task := agentTask(taskID, request.OwnerSessionID, database.TaskRunning, "") + + fake.tasks[taskID] = task + if fake.submitCh != nil { + fake.submitCh <- taskID + } + + return cloneAgentTask(task), nil +} + +func (fake *fakeController) Get(_ context.Context, taskID string) (*database.AgentTaskEntity, bool, error) { + fake.mu.Lock() + defer fake.mu.Unlock() + + task, found := fake.tasks[taskID] + if !found { + return nil, false, nil + } + + return cloneAgentTask(task), true, nil +} + +func (fake *fakeController) Await(ctx context.Context, taskID string) (*database.AgentTaskEntity, error) { + if fake.await != nil { + return fake.await(ctx, taskID) + } + + fake.mu.Lock() + defer fake.mu.Unlock() + + task := fake.tasks[taskID] + task.Task.State = database.TaskSucceeded + task.Task.Result = "done" + + return cloneAgentTask(task), nil +} + +func (fake *fakeController) Cancel(_ context.Context, owner, taskID string) (*database.TaskEntity, bool, error) { + fake.mu.Lock() + defer fake.mu.Unlock() + + fake.cancels = append(fake.cancels, [2]string{owner, taskID}) + task := fake.tasks[taskID] + task.Task.State = database.TaskCanceled + + result := task.Task + + return &result, true, nil +} + +func cloneAgentTask(task *database.AgentTaskEntity) *database.AgentTaskEntity { + clone := *task + + return &clone +} + +func TestRunnerSingleTask(t *testing.T) { + t.Parallel() + + fake := newFakeController() + runner, err := workflow.NewRunner(fake) + require.NoError(t, err) + + var events []workflow.Event + + result, err := runner.Run(context.Background(), &workflow.RunRequest{ + Arguments: nil, RunID: "", OnEvent: func(_ context.Context, event workflow.Event) error { + events = append(events, event) + + return nil + }, + Name: "", Source: agentSource, OwnerSessionID: testOwner, + PersistedLinks: nil, + }) + require.NoError(t, err) + + assert.Equal(t, []string{firstTask}, result.LaunchedTaskIDs) + require.Len(t, fake.submits, 1) + assert.Equal(t, testOwner, fake.submits[0].OwnerSessionID) + assert.Equal(t, "inspect", fake.submits[0].Prompt) + require.Len(t, events, 2) + assert.Equal(t, workflow.EventTaskLaunched, events[0].Kind) + assert.Equal(t, workflow.EventTaskCompleted, events[1].Kind) + assert.Equal(t, "done", events[1].Task.Result) +} + +func TestRunnerReusesPersistedInvocationByNormalizedNodeKey(t *testing.T) { + t.Parallel() + + fake := newFakeController() + fake.tasks[firstTask] = agentTask(firstTask, testOwner, database.TaskSucceeded, "persisted") + runner, err := workflow.NewRunner(fake) + require.NoError(t, err) + + const source = `import "librecode/workflow"; id, _ := workflow.Agent("inspect"); workflow.Wait(id)` + + result, err := runner.Run(t.Context(), &workflow.RunRequest{ + RunID: "", Name: "", Source: source, OwnerSessionID: testOwner, + PersistedLinks: []database.WorkflowAgentTaskEntity{{ + CreatedAt: time.Time{}, WorkflowTaskID: "", AgentTaskID: firstTask, + NodeKey: " agent ", InvocationIndex: 0, Sequence: 1, + }}, + Arguments: nil, OnEvent: nil, + }) + require.NoError(t, err) + assert.Empty(t, fake.submits) + assert.Equal(t, []string{firstTask}, result.LaunchedTaskIDs) + assert.Equal(t, "done", result.TaskResults[0].Result) +} + +func TestRunnerReusesPersistedInvocationsByNormalizedNodeKeyAndIndex(t *testing.T) { + t.Parallel() + + fake := newFakeController() + fake.tasks[firstTask] = agentTask(firstTask, testOwner, database.TaskSucceeded, "first") + fake.tasks[firstTask].Prompt = "does not match first call" + fake.tasks[secondTask] = agentTask(secondTask, testOwner, database.TaskSucceeded, "second") + fake.tasks[secondTask].Prompt = "does not match second call" + runner, err := workflow.NewRunner(fake) + require.NoError(t, err) + + const source = `import "librecode/workflow" + first, _ := workflow.Agent("alpha", map[string]any{"node_key": "review"}) + second, _ := workflow.Agent("beta", map[string]any{"node_key": " review "}) + []any{first, second}` + + result, err := runner.Run(t.Context(), &workflow.RunRequest{ + RunID: "", Name: "", Source: source, OwnerSessionID: testOwner, OnEvent: nil, Arguments: nil, + PersistedLinks: []database.WorkflowAgentTaskEntity{ + { + CreatedAt: time.Time{}, WorkflowTaskID: "", AgentTaskID: secondTask, + NodeKey: " review ", InvocationIndex: 1, Sequence: 0, + }, + { + CreatedAt: time.Time{}, WorkflowTaskID: "", AgentTaskID: firstTask, + NodeKey: "review", InvocationIndex: 0, Sequence: 0, + }, + }, + }) + require.NoError(t, err) + assert.Empty(t, fake.submits) + assert.Equal(t, []any{firstTask, secondTask}, result.Value) +} + +func TestRunnerScopesWaitToLaunchedOwnedTasks(t *testing.T) { + t.Parallel() + + tests := []struct { + setup func(*fakeController) + source string + name string + wantCancels [][2]string + }{ + { + name: "unlaunched task", + source: `import "librecode/workflow"; workflow.Wait("other")`, + wantCancels: nil, + setup: func(fake *fakeController) { + fake.tasks["other"] = agentTask("other", testOwner, database.TaskSucceeded, "secret") + }, + }, + { + name: "controller owner mismatch", + source: agentSource, + wantCancels: [][2]string{{testOwner, firstTask}}, + setup: func(fake *fakeController) { + fake.await = func(_ context.Context, taskID string) (*database.AgentTaskEntity, error) { + return agentTask(taskID, "other-session", database.TaskSucceeded, "secret"), nil + } + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + fake := newFakeController() + test.setup(fake) + runner, err := workflow.NewRunner(fake) + require.NoError(t, err) + + _, err = runner.Run(context.Background(), &workflow.RunRequest{ + Arguments: nil, RunID: "", OnEvent: nil, Name: "", Source: test.source, OwnerSessionID: testOwner, + PersistedLinks: nil, + }) + require.Error(t, err) + assert.Equal(t, test.wantCancels, fake.cancels) + }) + } +} + +func TestRunnerListAndCancelStayScopedToLaunchOrder(t *testing.T) { + t.Parallel() + + fake := newFakeController() + runner, err := workflow.NewRunner(fake) + require.NoError(t, err) + + const source = `import "librecode/workflow" +first, _ := workflow.Agent("first") +second, _ := workflow.Agent("second") +before, _ := workflow.List() +canceled, _ := workflow.Cancel(second) +[]any{first, second, before, canceled}` + + result, err := runner.Run(context.Background(), &workflow.RunRequest{ + Arguments: nil, RunID: "", OnEvent: nil, Name: "", Source: source, OwnerSessionID: testOwner, + PersistedLinks: nil, + }) + require.NoError(t, err) + + assert.Equal(t, []string{firstTask, secondTask}, result.LaunchedTaskIDs) + require.Len(t, result.TaskResults, 2) + assert.Equal(t, firstTask, result.TaskResults[0].ID) + assert.Equal(t, secondTask, result.TaskResults[1].ID) + assert.Equal(t, string(database.TaskCanceled), result.TaskResults[1].State) + assert.Equal(t, [][2]string{{testOwner, secondTask}}, fake.cancels) +} + +func TestRunnerPipelinePreservesInputOrder(t *testing.T) { + t.Parallel() + + fake := newFakeController() + runner, err := workflow.NewRunner(fake) + require.NoError(t, err) + + const source = `import "librecode/workflow" +results, _ := workflow.Pipeline([]any{1, 2, 3}, func(item any) (any, error) { + value := item.(int) + return value * value, nil +}, 2) +results` + + result, err := runner.Run(context.Background(), &workflow.RunRequest{ + Arguments: nil, RunID: "", OnEvent: nil, Name: "", Source: source, OwnerSessionID: testOwner, + PersistedLinks: nil, + }) + require.NoError(t, err) + + assert.Equal(t, []workflow.PipelineResult{ + {Index: 0, Value: 1, Error: ""}, + {Index: 1, Value: 4, Error: ""}, + {Index: 2, Value: 9, Error: ""}, + }, result.Value) +} + +func TestRunnerPipelineHandlesEmptyInput(t *testing.T) { + t.Parallel() + + fake := newFakeController() + runner, err := workflow.NewRunner(fake) + require.NoError(t, err) + + const source = `import "librecode/workflow" +results, _ := workflow.Pipeline([]any{}, func(item any) (any, error) { return item, nil }, 4) +results` + + result, err := runner.Run(context.Background(), &workflow.RunRequest{ + Arguments: nil, RunID: "", OnEvent: nil, Name: "", Source: source, OwnerSessionID: testOwner, + PersistedLinks: nil, + }) + require.NoError(t, err) + assert.Equal(t, []workflow.PipelineResult{}, result.Value) +} + +func TestRunnerDoesNotInferPipelineResultFromValueShape(t *testing.T) { + t.Parallel() + + runner, err := workflow.NewRunner(newFakeController()) + require.NoError(t, err) + + result, err := runner.Run(t.Context(), &workflow.RunRequest{ + RunID: "", Name: "", Source: `[]map[string]any{{"index": 0, "value": "ordinary", "error": ""}}`, + OwnerSessionID: testOwner, OnEvent: nil, Arguments: nil, PersistedLinks: nil, + }) + require.NoError(t, err) + + assert.IsType(t, []any{}, result.Value) + assert.NotEqual(t, []workflow.PipelineResult{{Value: "ordinary", Error: "", Index: 0}}, result.Value) +} + +func TestRunnerPipelineBoundsConcurrency(t *testing.T) { + t.Parallel() + + fake := newFakeController() + + var ( + active atomic.Int64 + peak atomic.Int64 + ) + + fake.await = func(_ context.Context, taskID string) (*database.AgentTaskEntity, error) { + current := active.Add(1) + updatePeak(&peak, current) + + time.Sleep(10 * time.Millisecond) + active.Add(-1) + + return agentTask(taskID, testOwner, database.TaskSucceeded, taskID), nil + } + + runner, err := workflow.NewRunner(fake) + require.NoError(t, err) + + const source = `import "librecode/workflow" +results, _ := workflow.Pipeline([]any{"a", "b", "c", "d"}, func(item any) (any, error) { + id, err := workflow.Agent(item.(string)) + if err != nil { return nil, err } + return workflow.Wait(id) +}, 2) +results` + + result, err := runner.Run(context.Background(), &workflow.RunRequest{ + Arguments: nil, RunID: "", OnEvent: nil, Name: "", Source: source, OwnerSessionID: testOwner, + PersistedLinks: nil, + }) + require.NoError(t, err) + require.Len(t, result.Value, 4) + assert.Positive(t, peak.Load()) + assert.LessOrEqual(t, peak.Load(), int64(2)) +} + +func TestRunnerPipelineCollectsFailureAndStopsScheduling(t *testing.T) { + t.Parallel() + + fake := newFakeController() + runner, err := workflow.NewRunner(fake) + require.NoError(t, err) + + const source = `import "librecode/workflow" +results, _ := workflow.Pipeline([]any{1, 2, 3}, func(item any) (any, error) { + value := item.(int) + if value == 1 { return workflow.Wait("other") } + return value, nil +}, 1) +results` + + result, err := runner.Run(context.Background(), &workflow.RunRequest{ + Arguments: nil, RunID: "", OnEvent: nil, Name: "", Source: source, OwnerSessionID: testOwner, + PersistedLinks: nil, + }) + require.NoError(t, err) + assert.Equal(t, []workflow.PipelineResult{ + { + Index: 0, + Value: workflow.TaskResult{ + ID: "", State: "", Result: "", ErrorCode: "", ErrorMessage: "", + }, + Error: "task was not launched by this workflow", + }, + {Index: 1, Value: nil, Error: "pipeline stopped before item was scheduled"}, + {Index: 2, Value: nil, Error: "pipeline stopped before item was scheduled"}, + }, result.Value) +} + +func TestRunnerPipelineRejectsInvalidConcurrency(t *testing.T) { + t.Parallel() + + fake := newFakeController() + runner, err := workflow.NewRunner(fake) + require.NoError(t, err) + + const source = `import "librecode/workflow" +_, err := workflow.Pipeline([]any{1}, func(item any) (any, error) { return item, nil }, 0) +err` + + _, err = runner.Run(context.Background(), &workflow.RunRequest{ + Arguments: nil, RunID: "", OnEvent: nil, Name: "", Source: source, OwnerSessionID: testOwner, + PersistedLinks: nil, + }) + require.Error(t, err) + assert.ErrorContains(t, err, "pipeline concurrency must be positive") +} + +func TestRunnerRejectsCancelForUnlaunchedTask(t *testing.T) { + t.Parallel() + + fake := newFakeController() + fake.tasks["other"] = agentTask("other", testOwner, database.TaskRunning, "") + runner, err := workflow.NewRunner(fake) + require.NoError(t, err) + + _, err = runner.Run(context.Background(), &workflow.RunRequest{ + Arguments: nil, RunID: "", OnEvent: nil, Name: "", + Source: `import "librecode/workflow"; workflow.Cancel("other")`, + OwnerSessionID: testOwner, + PersistedLinks: nil, + }) + require.Error(t, err) + assert.Empty(t, fake.cancels) +} + +func TestRunnerEventFailureCancelsActiveLaunchedTasks(t *testing.T) { + t.Parallel() + + fake := newFakeController() + runner, err := workflow.NewRunner(fake) + require.NoError(t, err) + + result, err := runner.Run(context.Background(), &workflow.RunRequest{ + Arguments: nil, RunID: "", OnEvent: func(_ context.Context, _ workflow.Event) error { + return assert.AnError + }, + Name: "", Source: `import "librecode/workflow"; workflow.Agent("inspect")`, + OwnerSessionID: testOwner, + PersistedLinks: nil, + }) + require.Error(t, err) + assert.Equal(t, [][2]string{{testOwner, firstTask}}, fake.cancels) + require.Len(t, result.TaskResults, 1) + assert.Equal(t, string(database.TaskCanceled), result.TaskResults[0].State) +} + +func TestRunnerCancellationCancelsActiveLaunchedTasks(t *testing.T) { + t.Parallel() + + fake := newFakeController() + fake.submitCh = make(chan string, 1) + fake.await = func(ctx context.Context, _ string) (*database.AgentTaskEntity, error) { + <-ctx.Done() + + return nil, ctx.Err() + } + runner, err := workflow.NewRunner(fake) + require.NoError(t, err) + + ctx, cancel := context.WithCancel(context.Background()) + + done := make(chan error, 1) + + go func() { + _, runErr := runner.Run(ctx, &workflow.RunRequest{ + Arguments: nil, RunID: "", OnEvent: nil, Name: "", Source: agentSource, OwnerSessionID: testOwner, + PersistedLinks: nil, + }) + done <- runErr + }() + + select { + case <-fake.submitCh: + cancel() + case <-time.After(time.Second): + t.Fatal("workflow did not submit task") + } + + require.Error(t, <-done) + assert.Equal(t, [][2]string{{testOwner, firstTask}}, fake.cancels) +} + +func updatePeak(peak *atomic.Int64, value int64) { + for { + current := peak.Load() + if value <= current || peak.CompareAndSwap(current, value) { + return + } + } +} + +func newFakeController() *fakeController { + return &fakeController{ + tasks: make(map[string]*database.AgentTaskEntity), await: nil, submitCh: nil, taskIDs: nil, + submits: nil, cancels: nil, mu: sync.Mutex{}, + } +} + +func agentTask(id, owner string, state database.TaskState, result string) *database.AgentTaskEntity { + return &database.AgentTaskEntity{ + Task: database.TaskEntity{ + CreatedAt: time.Time{}, StartedAt: nil, FinishedAt: nil, UpdatedAt: time.Time{}, LeaseExpiresAt: nil, + ID: id, Kind: database.TaskKindAgent, ParentTaskID: "", OwnerSessionID: owner, + ConcurrencyKey: "", LeaseOwner: "", State: state, Result: result, ErrorCode: "", ErrorMessage: "", + }, + ChildSessionID: "", AgentName: "", Prompt: "", Model: "", Provider: "", PolicyJSON: "", + UsageJSON: "", Depth: 0, + } +}