diff --git a/cmd/ask.go b/cmd/ask.go index 0024f1a2..0871380f 100644 --- a/cmd/ask.go +++ b/cmd/ask.go @@ -148,6 +148,15 @@ func init() { rootCmd.AddCommand(askCmd) } +func persistAskAssistantResponse(ctx context.Context, store session.Store, sess *session.Session, assistantMsg llm.Message, durationMs int64, reasoningCfg config.ReasoningConfig) error { + sessionMsg := session.NewMessageWithReasoningPolicy(sess.ID, assistantMsg, -1, reasoningCfg) + sessionMsg.DurationMs = durationMs + if err := store.AddMessage(ctx, sess.ID, sessionMsg); err != nil { + return fmt.Errorf("persist assistant response: %w", err) + } + return nil +} + func runAsk(cmd *cobra.Command, args []string) error { // Extract @agent from args if present atAgent, filteredArgs := ExtractAgentFromArgs(args) @@ -605,13 +614,7 @@ func runAsk(cmd *cobra.Command, args []string) error { // Set up response callback to save assistant message immediately (before tool execution) // This ensures the message is persisted even if tool execution fails/crashes persistResponseCompleted = func(ctx context.Context, turnIndex int, assistantMsg llm.Message, metrics llm.TurnMetrics) error { - // Calculate duration from stream start - durationMs := time.Since(turnStartTime).Milliseconds() - - sessionMsg := session.NewMessageWithReasoningPolicy(sess.ID, assistantMsg, -1, reasoningPersistenceCfg) - sessionMsg.DurationMs = durationMs - _ = store.AddMessage(ctx, sess.ID, sessionMsg) - return nil + return persistAskAssistantResponse(ctx, store, sess, assistantMsg, time.Since(turnStartTime).Milliseconds(), reasoningPersistenceCfg) } // Set up turn callback for tool result messages and metrics diff --git a/cmd/ask_test.go b/cmd/ask_test.go index 3bf0af26..d58388a9 100644 --- a/cmd/ask_test.go +++ b/cmd/ask_test.go @@ -4,6 +4,8 @@ import ( "bytes" "context" "encoding/json" + "errors" + "io" "strings" "testing" "time" @@ -11,6 +13,7 @@ import ( tea "charm.land/bubbletea/v2" "github.com/samsaffron/term-llm/internal/config" "github.com/samsaffron/term-llm/internal/llm" + "github.com/samsaffron/term-llm/internal/session" "github.com/samsaffron/term-llm/internal/ui" ) @@ -33,6 +36,86 @@ Some nice syntax improvements: It's an exciting release that balances new experimental features with practical improvements!` +type askAssistantFailOnceStore struct { + session.NoopStore + failErr error + failAssistantOnce bool + failedAssistant bool + messages []session.Message +} + +func (s *askAssistantFailOnceStore) AddMessage(ctx context.Context, sessionID string, msg *session.Message) error { + if s.failAssistantOnce && msg.Role == llm.RoleAssistant { + s.failAssistantOnce = false + s.failedAssistant = true + return s.failErr + } + s.messages = append(s.messages, *msg) + return nil +} + +func TestAskPersistenceRetriesAssistantThroughTurnCallbackWhenResponsePersistenceFails(t *testing.T) { + persistErr := errors.New("sqlite busy once") + store := &askAssistantFailOnceStore{failErr: persistErr, failAssistantOnce: true} + sess := &session.Session{ID: "ask-persist-retry"} + reasoningCfg := config.DefaultReasoningConfig() + + tool := &progressiveTestTool{} + registry := llm.NewToolRegistry() + registry.Register(tool) + provider := llm.NewMockProvider("mock").WithCapabilities(llm.Capabilities{ToolCalls: true}) + provider.AddToolCall("call-1", "progressive_test_tool", map[string]any{}) + engine := llm.NewEngine(provider, registry) + + engine.SetResponseCompletedCallback(func(ctx context.Context, turnIndex int, assistantMsg llm.Message, metrics llm.TurnMetrics) error { + return persistAskAssistantResponse(ctx, store, sess, assistantMsg, 123, reasoningCfg) + }) + engine.SetTurnCompletedCallback(func(ctx context.Context, turnIndex int, messages []llm.Message, metrics llm.TurnMetrics) error { + for _, msg := range messages { + sessionMsg := session.NewMessageWithReasoningPolicy(sess.ID, msg, -1, reasoningCfg) + if msg.Role == llm.RoleAssistant { + sessionMsg.DurationMs = 123 + } + _ = store.AddMessage(ctx, sess.ID, sessionMsg) + } + return nil + }) + + stream, err := engine.Stream(context.Background(), llm.Request{ + Messages: []llm.Message{llm.UserText("use the tool")}, + Tools: []llm.ToolSpec{tool.Spec()}, + }) + if err != nil { + t.Fatalf("stream error: %v", err) + } + defer stream.Close() + for { + _, err := stream.Recv() + if err == io.EOF { + break + } + if err != nil { + t.Fatalf("recv error: %v", err) + } + } + + if !store.failedAssistant { + t.Fatal("expected response AddMessage to fail once") + } + if len(store.messages) != 2 { + t.Fatalf("persisted messages = %d, want assistant + tool result via turn callback", len(store.messages)) + } + if store.messages[0].Role != llm.RoleAssistant { + t.Fatalf("first persisted role = %s, want assistant", store.messages[0].Role) + } + if store.messages[0].DurationMs != 123 { + t.Fatalf("assistant duration = %d, want 123", store.messages[0].DurationMs) + } + if store.messages[1].Role != llm.RoleTool { + t.Fatalf("second persisted role = %s, want tool", store.messages[1].Role) + } +} + func TestRunOnCompleteCapture_CapturesStdout(t *testing.T) { result, err := runOnCompleteCapture("cat", "hello") if err != nil {