feat(backend): add thinking/reasoning streaming support (#113)
This commit is contained in:
@@ -357,6 +357,28 @@ func (a *messageCreatorAdapter) AppendAgentMessage(ctx context.Context, messageI
|
||||
return a.svc.AppendMessageContent(ctx, messageID, additionalContent)
|
||||
}
|
||||
|
||||
// CreateThinkingMessageStreaming creates a new thinking message with a pre-generated ID.
|
||||
// This is used for real-time streaming of agent thinking/reasoning content.
|
||||
func (a *messageCreatorAdapter) CreateThinkingMessageStreaming(ctx context.Context, messageID, taskID, content, agentSessionID, turnID string) error {
|
||||
_, err := a.svc.CreateMessageWithID(ctx, messageID, &taskservice.CreateMessageRequest{
|
||||
TaskSessionID: agentSessionID,
|
||||
TaskID: taskID,
|
||||
TurnID: turnID,
|
||||
Content: "",
|
||||
AuthorType: "agent",
|
||||
Type: "thinking",
|
||||
Metadata: map[string]interface{}{
|
||||
"thinking": content,
|
||||
},
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
// AppendThinkingMessage appends additional content to an existing streaming thinking message.
|
||||
func (a *messageCreatorAdapter) AppendThinkingMessage(ctx context.Context, messageID, additionalContent string) error {
|
||||
return a.svc.AppendThinkingContent(ctx, messageID, additionalContent)
|
||||
}
|
||||
|
||||
// turnServiceAdapter adapts the task service to the orchestrator.TurnService interface
|
||||
type turnServiceAdapter struct {
|
||||
svc *taskservice.Service
|
||||
|
||||
@@ -48,11 +48,13 @@ type AgentStreamEventData struct {
|
||||
SessionStatus string `json:"session_status,omitempty"` // "resumed" or "new" for session_status events
|
||||
Data interface{} `json:"data,omitempty"`
|
||||
|
||||
// Streaming message fields (for "message_streaming" type)
|
||||
// Streaming message fields (for "message_streaming" and "thinking_streaming" types)
|
||||
// MessageID is the ID of the message being streamed (empty for first chunk, set for appends)
|
||||
MessageID string `json:"message_id,omitempty"`
|
||||
// IsAppend indicates whether this is an append to an existing message (true) or a new message (false)
|
||||
IsAppend bool `json:"is_append,omitempty"`
|
||||
// MessageType distinguishes between "message" and "thinking" content types
|
||||
MessageType string `json:"message_type,omitempty"`
|
||||
}
|
||||
|
||||
// AgentStreamEventPayload is the payload for agent stream events (WebSocket streaming).
|
||||
|
||||
@@ -1236,15 +1236,33 @@ func (m *Manager) handleAgentEvent(execution *AgentExecution, event agentctl.Age
|
||||
}
|
||||
|
||||
case "reasoning":
|
||||
// Accumulate reasoning/thinking content
|
||||
execution.messageMu.Lock()
|
||||
// Stream thinking content like message chunks for real-time feedback
|
||||
// All adapters normalize reasoning to ReasoningText field
|
||||
if event.ReasoningText != "" {
|
||||
execution.reasoningBuffer.WriteString(event.ReasoningText)
|
||||
execution.messageMu.Lock()
|
||||
execution.thinkingBuffer.WriteString(event.ReasoningText)
|
||||
|
||||
// Check if buffer contains newlines - if so, flush up to the last newline
|
||||
bufContent := execution.thinkingBuffer.String()
|
||||
lastNewline := strings.LastIndex(bufContent, "\n")
|
||||
if lastNewline != -1 {
|
||||
// Extract content up to and including the last newline
|
||||
toFlush := bufContent[:lastNewline+1]
|
||||
remainder := bufContent[lastNewline+1:]
|
||||
|
||||
// Reset buffer with remainder
|
||||
execution.thinkingBuffer.Reset()
|
||||
execution.thinkingBuffer.WriteString(remainder)
|
||||
execution.messageMu.Unlock()
|
||||
|
||||
// Publish streaming thinking event if there's content to flush
|
||||
if strings.TrimSpace(toFlush) != "" {
|
||||
m.publishStreamingThinking(execution, toFlush)
|
||||
}
|
||||
} else {
|
||||
execution.messageMu.Unlock()
|
||||
}
|
||||
}
|
||||
if event.ReasoningSummary != "" {
|
||||
execution.summaryBuffer.WriteString(event.ReasoningSummary)
|
||||
}
|
||||
execution.messageMu.Unlock()
|
||||
|
||||
case "tool_call":
|
||||
// Tool call starting marks a step boundary - flush the accumulated message
|
||||
@@ -1459,17 +1477,33 @@ func (m *Manager) publishStreamingMessage(execution *AgentExecution, content str
|
||||
// flushMessageBuffer extracts any accumulated message from the buffer and returns it.
|
||||
// This is called when a tool use starts or on complete to get the agent's response.
|
||||
// It also clears the currentMessageID to start fresh for the next message segment.
|
||||
// Additionally flushes any accumulated thinking content.
|
||||
func (m *Manager) flushMessageBuffer(execution *AgentExecution) string {
|
||||
execution.messageMu.Lock()
|
||||
agentMessage := execution.messageBuffer.String()
|
||||
thinkingContent := execution.thinkingBuffer.String()
|
||||
execution.messageBuffer.Reset()
|
||||
execution.reasoningBuffer.Reset()
|
||||
execution.summaryBuffer.Reset()
|
||||
// Clear the streaming message ID so next segment starts fresh
|
||||
execution.thinkingBuffer.Reset()
|
||||
// Clear the streaming message IDs so next segment starts fresh
|
||||
currentMsgID := execution.currentMessageID
|
||||
currentThinkingID := execution.currentThinkingID
|
||||
execution.currentMessageID = ""
|
||||
execution.currentThinkingID = ""
|
||||
execution.messageMu.Unlock()
|
||||
|
||||
// If we have remaining thinking content, publish it
|
||||
trimmedThinking := strings.TrimSpace(thinkingContent)
|
||||
if trimmedThinking != "" {
|
||||
if currentThinkingID != "" {
|
||||
// Append to existing streaming thinking message
|
||||
m.publishStreamingThinkingFinal(execution, currentThinkingID, trimmedThinking)
|
||||
} else {
|
||||
// No streaming thinking message exists yet - create one with all the content
|
||||
// This happens when thinking content has no newlines (never triggered streaming)
|
||||
m.publishStreamingThinking(execution, trimmedThinking)
|
||||
}
|
||||
}
|
||||
|
||||
// If we have remaining content and an active streaming message, append it
|
||||
trimmedMessage := strings.TrimSpace(agentMessage)
|
||||
if trimmedMessage != "" && currentMsgID != "" {
|
||||
@@ -1509,6 +1543,71 @@ func (m *Manager) publishStreamingMessageFinal(execution *AgentExecution, messag
|
||||
m.eventPublisher.PublishAgentStreamEventPayload(payload)
|
||||
}
|
||||
|
||||
// publishStreamingThinking publishes a streaming thinking event for real-time thinking updates.
|
||||
// It creates a new thinking message on first call (currentThinkingID empty) or appends to existing.
|
||||
// The message ID is generated and set synchronously to avoid race conditions.
|
||||
func (m *Manager) publishStreamingThinking(execution *AgentExecution, content string) {
|
||||
execution.messageMu.Lock()
|
||||
isAppend := execution.currentThinkingID != ""
|
||||
thinkingID := execution.currentThinkingID
|
||||
|
||||
// If this is the first chunk of a new thinking segment, generate the ID now
|
||||
if !isAppend {
|
||||
thinkingID = uuid.New().String()
|
||||
execution.currentThinkingID = thinkingID
|
||||
}
|
||||
execution.messageMu.Unlock()
|
||||
|
||||
event := AgentStreamEventData{
|
||||
Type: "thinking_streaming",
|
||||
Text: content,
|
||||
MessageID: thinkingID,
|
||||
IsAppend: isAppend,
|
||||
MessageType: "thinking",
|
||||
}
|
||||
|
||||
// Create payload manually to include streaming-specific fields
|
||||
payload := &AgentStreamEventPayload{
|
||||
Type: "agent/event",
|
||||
Timestamp: time.Now().UTC().Format(time.RFC3339Nano),
|
||||
AgentID: execution.ID,
|
||||
TaskID: execution.TaskID,
|
||||
SessionID: execution.SessionID,
|
||||
Data: &event,
|
||||
}
|
||||
|
||||
// Publish the streaming event - orchestrator will handle create/append logic
|
||||
m.eventPublisher.PublishAgentStreamEventPayload(payload)
|
||||
}
|
||||
|
||||
// publishStreamingThinkingFinal publishes the final chunk of a streaming thinking message.
|
||||
// This is called during flush to append any remaining buffered thinking content.
|
||||
func (m *Manager) publishStreamingThinkingFinal(execution *AgentExecution, thinkingID, content string) {
|
||||
event := AgentStreamEventData{
|
||||
Type: "thinking_streaming",
|
||||
Text: content,
|
||||
MessageID: thinkingID,
|
||||
IsAppend: true,
|
||||
MessageType: "thinking",
|
||||
}
|
||||
|
||||
payload := &AgentStreamEventPayload{
|
||||
Type: "agent/event",
|
||||
Timestamp: time.Now().UTC().Format(time.RFC3339Nano),
|
||||
AgentID: execution.ID,
|
||||
TaskID: execution.TaskID,
|
||||
SessionID: execution.SessionID,
|
||||
Data: &event,
|
||||
}
|
||||
|
||||
m.logger.Debug("publishing final streaming thinking chunk",
|
||||
zap.String("execution_id", execution.ID),
|
||||
zap.String("thinking_id", thinkingID),
|
||||
zap.Int("content_length", len(content)))
|
||||
|
||||
m.eventPublisher.PublishAgentStreamEventPayload(payload)
|
||||
}
|
||||
|
||||
// updateExecutionError updates an execution with an error
|
||||
func (m *Manager) updateExecutionError(executionID, errorMsg string) {
|
||||
m.executionStore.UpdateError(executionID, errorMsg)
|
||||
@@ -1550,9 +1649,9 @@ func (m *Manager) CancelAgent(ctx context.Context, executionID string) error {
|
||||
// Clear streaming state after cancel to ensure clean state for next prompt
|
||||
execution.messageMu.Lock()
|
||||
execution.messageBuffer.Reset()
|
||||
execution.reasoningBuffer.Reset()
|
||||
execution.summaryBuffer.Reset()
|
||||
execution.thinkingBuffer.Reset()
|
||||
execution.currentMessageID = ""
|
||||
execution.currentThinkingID = ""
|
||||
execution.messageMu.Unlock()
|
||||
|
||||
// Mark as ready for follow-up prompts after successful cancel
|
||||
|
||||
@@ -384,9 +384,9 @@ func (sm *SessionManager) SendPrompt(
|
||||
// This ensures each prompt starts fresh and doesn't append to previous message
|
||||
execution.messageMu.Lock()
|
||||
execution.messageBuffer.Reset()
|
||||
execution.reasoningBuffer.Reset()
|
||||
execution.summaryBuffer.Reset()
|
||||
execution.currentMessageID = "" // Clear streaming message ID for new turn
|
||||
execution.thinkingBuffer.Reset()
|
||||
execution.currentMessageID = "" // Clear streaming message ID for new turn
|
||||
execution.currentThinkingID = "" // Clear streaming thinking ID for new turn
|
||||
execution.messageMu.Unlock()
|
||||
|
||||
// Check if we need to inject resume context (fork_session pattern)
|
||||
@@ -439,10 +439,10 @@ func (sm *SessionManager) SendPrompt(
|
||||
execution.messageMu.Lock()
|
||||
agentMessage := execution.messageBuffer.String()
|
||||
execution.messageBuffer.Reset()
|
||||
execution.reasoningBuffer.Reset()
|
||||
execution.summaryBuffer.Reset()
|
||||
execution.thinkingBuffer.Reset()
|
||||
currentMsgID := execution.currentMessageID
|
||||
execution.currentMessageID = "" // Clear for next turn
|
||||
execution.currentMessageID = "" // Clear for next turn
|
||||
execution.currentThinkingID = "" // Clear for next turn
|
||||
execution.messageMu.Unlock()
|
||||
|
||||
stopReason := ""
|
||||
|
||||
@@ -50,14 +50,14 @@ type AgentExecution struct {
|
||||
PassthroughProcessID string // Process ID in the interactive runner (empty if not in passthrough mode)
|
||||
|
||||
// Buffers for accumulating agent response during a prompt
|
||||
messageBuffer strings.Builder
|
||||
reasoningBuffer strings.Builder
|
||||
summaryBuffer strings.Builder
|
||||
messageMu sync.Mutex
|
||||
messageBuffer strings.Builder
|
||||
thinkingBuffer strings.Builder
|
||||
messageMu sync.Mutex
|
||||
|
||||
// Streaming message tracking - ID of the current in-progress message being streamed
|
||||
// This is set when we create a streaming message and cleared on tool_call/complete
|
||||
currentMessageID string
|
||||
// Streaming message tracking - IDs of the current in-progress messages being streamed
|
||||
// These are set when we create a streaming message and cleared on tool_call/complete
|
||||
currentMessageID string
|
||||
currentThinkingID string
|
||||
|
||||
// Resume context tracking for fork_session pattern (ACP agents that don't support session/load)
|
||||
// needsResumeContext is set to true when the session has history that should be injected
|
||||
|
||||
@@ -253,26 +253,6 @@ func (c *Client) SessionUpdate(ctx context.Context, n acp.SessionNotification) e
|
||||
handler := c.updateHandler
|
||||
c.mu.RUnlock()
|
||||
|
||||
// Log the update type
|
||||
u := n.Update
|
||||
switch {
|
||||
case u.AgentMessageChunk != nil:
|
||||
if u.AgentMessageChunk.Content.Text != nil {
|
||||
c.logger.Debug("agent message chunk",
|
||||
zap.String("text", u.AgentMessageChunk.Content.Text.Text[:min(50, len(u.AgentMessageChunk.Content.Text.Text))]))
|
||||
}
|
||||
case u.ToolCall != nil:
|
||||
c.logger.Info("tool call",
|
||||
zap.String("tool_call_id", string(u.ToolCall.ToolCallId)),
|
||||
zap.String("title", u.ToolCall.Title),
|
||||
zap.String("status", string(u.ToolCall.Status)))
|
||||
case u.ToolCallUpdate != nil:
|
||||
c.logger.Debug("tool call update",
|
||||
zap.String("tool_call_id", string(u.ToolCallUpdate.ToolCallId)))
|
||||
case u.Plan != nil:
|
||||
c.logger.Info("plan update", zap.Int("entries", len(u.Plan.Entries)))
|
||||
}
|
||||
|
||||
// Forward to handler if set
|
||||
if handler != nil {
|
||||
handler(n)
|
||||
|
||||
@@ -342,14 +342,6 @@ func (a *ACPAdapter) handleACPUpdate(n acp.SessionNotification) {
|
||||
func (a *ACPAdapter) convertNotification(n acp.SessionNotification) *AgentEvent {
|
||||
u := n.Update
|
||||
|
||||
// Log what update types we're receiving
|
||||
a.logger.Debug("ACP notification received",
|
||||
zap.Bool("has_tool_call", u.ToolCall != nil),
|
||||
zap.Bool("has_tool_call_update", u.ToolCallUpdate != nil),
|
||||
zap.Bool("has_message_chunk", u.AgentMessageChunk != nil),
|
||||
zap.Bool("has_thought_chunk", u.AgentThoughtChunk != nil),
|
||||
zap.Bool("has_plan", u.Plan != nil))
|
||||
|
||||
switch {
|
||||
case u.AgentMessageChunk != nil:
|
||||
if u.AgentMessageChunk.Content.Text != nil {
|
||||
@@ -361,7 +353,8 @@ func (a *ACPAdapter) convertNotification(n acp.SessionNotification) *AgentEvent
|
||||
}
|
||||
|
||||
case u.AgentThoughtChunk != nil:
|
||||
// Agent thinking/reasoning - map to the new reasoning type
|
||||
// Agent thinking/reasoning - map to the reasoning type
|
||||
// Note: Only models with extended thinking (e.g., Opus 4.5) send agent_thought_chunk
|
||||
if u.AgentThoughtChunk.Content.Text != nil {
|
||||
return &AgentEvent{
|
||||
Type: EventTypeReasoning,
|
||||
|
||||
@@ -53,7 +53,6 @@ type CodexAdapter struct {
|
||||
// Accumulators for streaming content
|
||||
messageBuffer string
|
||||
reasoningBuffer string
|
||||
summaryBuffer string
|
||||
|
||||
// Turn completion signaling
|
||||
turnCompleteCh chan turnCompleteResult
|
||||
@@ -423,7 +422,6 @@ func (a *CodexAdapter) Prompt(ctx context.Context, message string) error {
|
||||
// Reset accumulators for new turn
|
||||
a.messageBuffer = ""
|
||||
a.reasoningBuffer = ""
|
||||
a.summaryBuffer = ""
|
||||
// Create channel to wait for turn completion
|
||||
a.turnCompleteCh = make(chan turnCompleteResult, 1)
|
||||
a.mu.Unlock()
|
||||
@@ -632,19 +630,21 @@ func (a *CodexAdapter) handleNotification(method string, params json.RawMessage)
|
||||
})
|
||||
|
||||
case codex.NotifyItemReasoningSummaryDelta:
|
||||
// Codex sends reasoning via summaryTextDelta - normalize to ReasoningText
|
||||
// so the lifecycle manager doesn't need to know about protocol differences
|
||||
var p codex.ReasoningDeltaParams
|
||||
if err := json.Unmarshal(params, &p); err != nil {
|
||||
a.logger.Warn("failed to parse reasoning summary delta", zap.Error(err))
|
||||
return
|
||||
}
|
||||
a.mu.Lock()
|
||||
a.summaryBuffer += p.Delta
|
||||
a.reasoningBuffer += p.Delta
|
||||
a.mu.Unlock()
|
||||
a.sendUpdate(AgentEvent{
|
||||
Type: EventTypeReasoning,
|
||||
SessionID: threadID,
|
||||
OperationID: turnID,
|
||||
ReasoningSummary: p.Delta,
|
||||
Type: EventTypeReasoning,
|
||||
SessionID: threadID,
|
||||
OperationID: turnID,
|
||||
ReasoningText: p.Delta,
|
||||
})
|
||||
|
||||
case codex.NotifyTurnCompleted:
|
||||
|
||||
@@ -88,7 +88,6 @@ type RepositoryDiscoveryConfig struct {
|
||||
type WorktreeConfig struct {
|
||||
Enabled bool `mapstructure:"enabled"` // Enable worktree mode
|
||||
BasePath string `mapstructure:"basePath"` // Base directory for worktrees (default: ~/.kandev/worktrees)
|
||||
MaxPerRepo int `mapstructure:"maxPerRepo"` // Max worktrees per repository (default: 10)
|
||||
DefaultBranch string `mapstructure:"defaultBranch"` // Default base branch (default: main)
|
||||
CleanupOnRemove bool `mapstructure:"cleanupOnRemove"` // Remove worktree directory on task deletion
|
||||
}
|
||||
@@ -204,7 +203,6 @@ func setDefaults(v *viper.Viper) {
|
||||
// Worktree defaults
|
||||
v.SetDefault("worktree.enabled", true)
|
||||
v.SetDefault("worktree.basePath", "~/.kandev/worktrees")
|
||||
v.SetDefault("worktree.maxPerRepo", 10)
|
||||
v.SetDefault("worktree.defaultBranch", "main")
|
||||
v.SetDefault("worktree.cleanupOnRemove", true)
|
||||
}
|
||||
|
||||
@@ -570,6 +570,25 @@ func (a *testMessageCreatorAdapter) AppendAgentMessage(ctx context.Context, mess
|
||||
return a.svc.AppendMessageContent(ctx, messageID, additionalContent)
|
||||
}
|
||||
|
||||
func (a *testMessageCreatorAdapter) CreateThinkingMessageStreaming(ctx context.Context, messageID, taskID, content, agentSessionID, turnID string) error {
|
||||
_, err := a.svc.CreateMessageWithID(ctx, messageID, &taskservice.CreateMessageRequest{
|
||||
TaskSessionID: agentSessionID,
|
||||
TaskID: taskID,
|
||||
TurnID: turnID,
|
||||
Content: "",
|
||||
AuthorType: "agent",
|
||||
Type: "thinking",
|
||||
Metadata: map[string]interface{}{
|
||||
"thinking": content,
|
||||
},
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func (a *testMessageCreatorAdapter) AppendThinkingMessage(ctx context.Context, messageID, additionalContent string) error {
|
||||
return a.svc.AppendThinkingContent(ctx, messageID, additionalContent)
|
||||
}
|
||||
|
||||
// testTurnServiceAdapter adapts the task service to the orchestrator.TurnService interface for tests
|
||||
type testTurnServiceAdapter struct {
|
||||
svc *taskservice.Service
|
||||
|
||||
@@ -0,0 +1,250 @@
|
||||
// Package integration provides end-to-end integration tests for the Kandev backend.
|
||||
// This file tests that agent thinking/reasoning events are properly saved to the database.
|
||||
package integration
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/kandev/kandev/internal/agent/lifecycle"
|
||||
"github.com/kandev/kandev/internal/events"
|
||||
"github.com/kandev/kandev/internal/events/bus"
|
||||
"github.com/kandev/kandev/internal/task/models"
|
||||
)
|
||||
|
||||
// TestThinkingEventsPersistToDatabase verifies that agent thinking/reasoning events
|
||||
// are properly saved to the database through the full orchestrator pipeline.
|
||||
//
|
||||
// This tests the flow:
|
||||
// 1. Lifecycle manager receives "reasoning" event from agentctl
|
||||
// 2. Manager buffers content and publishes "thinking_streaming" event
|
||||
// 3. Orchestrator handles event and calls messageCreator.CreateThinkingMessageStreaming
|
||||
// 4. Task service creates message with type="thinking" in database
|
||||
func TestThinkingEventsPersistToDatabase(t *testing.T) {
|
||||
ts := NewOrchestratorTestServer(t)
|
||||
defer ts.Close()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a task with session
|
||||
taskID := ts.CreateTestTask(t, "test-agent", 1)
|
||||
|
||||
// Create a session for the task using the repository directly
|
||||
sessionID := uuid.New().String()
|
||||
session := &models.TaskSession{
|
||||
ID: sessionID,
|
||||
TaskID: taskID,
|
||||
AgentProfileID: "test-agent",
|
||||
State: models.TaskSessionStateRunning,
|
||||
StartedAt: time.Now(),
|
||||
UpdatedAt: time.Now(),
|
||||
}
|
||||
err := ts.TaskRepo.CreateTaskSession(ctx, session)
|
||||
require.NoError(t, err, "failed to create session")
|
||||
|
||||
// Generate IDs for the test
|
||||
thinkingID := uuid.New().String()
|
||||
thinkingContent := "The user wants to list the files in the current directory.\nI'll use the view tool with type \"directory\" to show the contents of the workspace."
|
||||
|
||||
// Simulate a thinking_streaming event being published
|
||||
// This is what the lifecycle manager would publish after receiving a "reasoning" event
|
||||
eventData := &lifecycle.AgentStreamEventPayload{
|
||||
Type: "agent/event",
|
||||
Timestamp: time.Now().UTC().Format(time.RFC3339Nano),
|
||||
TaskID: taskID,
|
||||
SessionID: sessionID,
|
||||
Data: &lifecycle.AgentStreamEventData{
|
||||
Type: "thinking_streaming",
|
||||
Text: thinkingContent,
|
||||
MessageID: thinkingID,
|
||||
IsAppend: false, // First chunk - creates new message
|
||||
MessageType: "thinking",
|
||||
},
|
||||
}
|
||||
|
||||
// Publish the event to the event bus (simulating what lifecycle manager does)
|
||||
subject := events.BuildAgentStreamSubject(taskID)
|
||||
event := bus.NewEvent(events.AgentStream, "test", eventData)
|
||||
err = ts.EventBus.Publish(ctx, subject, event)
|
||||
require.NoError(t, err, "failed to publish thinking event")
|
||||
|
||||
// Wait for event to be processed
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
|
||||
// Check that the thinking message was saved to the database
|
||||
messages, err := ts.TaskSvc.ListMessages(ctx, sessionID)
|
||||
require.NoError(t, err, "failed to get messages")
|
||||
|
||||
// Find the thinking message
|
||||
var foundThinking bool
|
||||
for _, msg := range messages {
|
||||
t.Logf("Message: ID=%s, Type=%s, AuthorType=%s", msg.ID, msg.Type, msg.AuthorType)
|
||||
if msg.Type == "thinking" {
|
||||
foundThinking = true
|
||||
assert.Equal(t, thinkingID, msg.ID, "thinking message ID should match")
|
||||
assert.Equal(t, models.MessageAuthorAgent, msg.AuthorType, "thinking message author should be agent")
|
||||
// Check metadata contains thinking content
|
||||
thinking, hasThinking := msg.Metadata["thinking"]
|
||||
assert.True(t, hasThinking, "thinking message should have 'thinking' in metadata")
|
||||
assert.Equal(t, thinkingContent, thinking, "thinking content should match")
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
assert.True(t, foundThinking, "should find a thinking message in the database")
|
||||
}
|
||||
|
||||
// TestThinkingEventAppend verifies that appending to thinking messages works correctly.
|
||||
func TestThinkingEventAppend(t *testing.T) {
|
||||
ts := NewOrchestratorTestServer(t)
|
||||
defer ts.Close()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a task with session
|
||||
taskID := ts.CreateTestTask(t, "test-agent", 1)
|
||||
|
||||
// Create a session for the task using the repository directly
|
||||
sessionID := uuid.New().String()
|
||||
session := &models.TaskSession{
|
||||
ID: sessionID,
|
||||
TaskID: taskID,
|
||||
AgentProfileID: "test-agent",
|
||||
State: models.TaskSessionStateRunning,
|
||||
StartedAt: time.Now(),
|
||||
UpdatedAt: time.Now(),
|
||||
}
|
||||
err := ts.TaskRepo.CreateTaskSession(ctx, session)
|
||||
require.NoError(t, err)
|
||||
|
||||
thinkingID := uuid.New().String()
|
||||
chunk1 := "First chunk of thinking content.\n"
|
||||
chunk2 := "Second chunk with more reasoning."
|
||||
|
||||
// First chunk - creates the message
|
||||
event1 := &lifecycle.AgentStreamEventPayload{
|
||||
Type: "agent/event",
|
||||
Timestamp: time.Now().UTC().Format(time.RFC3339Nano),
|
||||
TaskID: taskID,
|
||||
SessionID: sessionID,
|
||||
Data: &lifecycle.AgentStreamEventData{
|
||||
Type: "thinking_streaming",
|
||||
Text: chunk1,
|
||||
MessageID: thinkingID,
|
||||
IsAppend: false,
|
||||
MessageType: "thinking",
|
||||
},
|
||||
}
|
||||
|
||||
subject := events.BuildAgentStreamSubject(taskID)
|
||||
err = ts.EventBus.Publish(ctx, subject, bus.NewEvent(events.AgentStream, "test", event1))
|
||||
require.NoError(t, err)
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
|
||||
// Second chunk - appends to existing message
|
||||
event2 := &lifecycle.AgentStreamEventPayload{
|
||||
Type: "agent/event",
|
||||
Timestamp: time.Now().UTC().Format(time.RFC3339Nano),
|
||||
TaskID: taskID,
|
||||
SessionID: sessionID,
|
||||
Data: &lifecycle.AgentStreamEventData{
|
||||
Type: "thinking_streaming",
|
||||
Text: chunk2,
|
||||
MessageID: thinkingID,
|
||||
IsAppend: true,
|
||||
MessageType: "thinking",
|
||||
},
|
||||
}
|
||||
|
||||
err = ts.EventBus.Publish(ctx, subject, bus.NewEvent(events.AgentStream, "test", event2))
|
||||
require.NoError(t, err)
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
|
||||
// Verify the message contains both chunks
|
||||
messages, err := ts.TaskSvc.ListMessages(ctx, sessionID)
|
||||
require.NoError(t, err)
|
||||
|
||||
var foundThinking bool
|
||||
for _, msg := range messages {
|
||||
if msg.Type == "thinking" && msg.ID == thinkingID {
|
||||
foundThinking = true
|
||||
thinking := msg.Metadata["thinking"].(string)
|
||||
assert.Contains(t, thinking, chunk1, "should contain first chunk")
|
||||
assert.Contains(t, thinking, chunk2, "should contain second chunk")
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
assert.True(t, foundThinking, "should find thinking message with both chunks")
|
||||
}
|
||||
|
||||
// TestCodexReasoningSummaryEvents verifies that Codex-style reasoning events
|
||||
// (which use ReasoningSummary instead of ReasoningText) are properly handled.
|
||||
// Codex sends item/reasoning/summaryTextDelta which sets ReasoningSummary.
|
||||
func TestCodexReasoningSummaryEvents(t *testing.T) {
|
||||
ts := NewOrchestratorTestServer(t)
|
||||
defer ts.Close()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a task with session
|
||||
taskID := ts.CreateTestTask(t, "test-agent", 1)
|
||||
|
||||
sessionID := uuid.New().String()
|
||||
session := &models.TaskSession{
|
||||
ID: sessionID,
|
||||
TaskID: taskID,
|
||||
AgentProfileID: "test-agent",
|
||||
State: models.TaskSessionStateRunning,
|
||||
StartedAt: time.Now(),
|
||||
UpdatedAt: time.Now(),
|
||||
}
|
||||
err := ts.TaskRepo.CreateTaskSession(ctx, session)
|
||||
require.NoError(t, err)
|
||||
|
||||
thinkingID := uuid.New().String()
|
||||
// Codex sends reasoning summaries like this
|
||||
summaryContent := "**Processing simple calculation**\nUser asks for a calculator response."
|
||||
|
||||
// Simulate a thinking_streaming event with content from ReasoningSummary
|
||||
// This is what would happen after lifecycle manager processes Codex's reasoning
|
||||
eventData := &lifecycle.AgentStreamEventPayload{
|
||||
Type: "agent/event",
|
||||
Timestamp: time.Now().UTC().Format(time.RFC3339Nano),
|
||||
TaskID: taskID,
|
||||
SessionID: sessionID,
|
||||
Data: &lifecycle.AgentStreamEventData{
|
||||
Type: "thinking_streaming",
|
||||
Text: summaryContent,
|
||||
MessageID: thinkingID,
|
||||
IsAppend: false,
|
||||
MessageType: "thinking",
|
||||
},
|
||||
}
|
||||
|
||||
subject := events.BuildAgentStreamSubject(taskID)
|
||||
err = ts.EventBus.Publish(ctx, subject, bus.NewEvent(events.AgentStream, "test", eventData))
|
||||
require.NoError(t, err)
|
||||
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
|
||||
messages, err := ts.TaskSvc.ListMessages(ctx, sessionID)
|
||||
require.NoError(t, err)
|
||||
|
||||
var foundThinking bool
|
||||
for _, msg := range messages {
|
||||
if msg.Type == "thinking" && msg.ID == thinkingID {
|
||||
foundThinking = true
|
||||
thinking := msg.Metadata["thinking"].(string)
|
||||
assert.Equal(t, summaryContent, thinking, "thinking content should match Codex summary")
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
assert.True(t, foundThinking, "should find thinking message from Codex-style event")
|
||||
}
|
||||
@@ -515,6 +515,9 @@ func (s *Service) handleAgentStreamEvent(ctx context.Context, payload *lifecycle
|
||||
case "message_streaming":
|
||||
s.handleMessageStreamingEvent(ctx, payload)
|
||||
|
||||
case "thinking_streaming":
|
||||
s.handleThinkingStreamingEvent(ctx, payload)
|
||||
|
||||
case "tool_call":
|
||||
s.saveAgentTextIfPresent(ctx, payload)
|
||||
s.handleToolCallEvent(ctx, payload)
|
||||
@@ -765,6 +768,56 @@ func (s *Service) handleMessageStreamingEvent(ctx context.Context, payload *life
|
||||
}
|
||||
}
|
||||
|
||||
// handleThinkingStreamingEvent handles streaming thinking events for real-time reasoning updates.
|
||||
// It creates a new thinking message on first chunk (IsAppend=false) or appends to existing (IsAppend=true).
|
||||
func (s *Service) handleThinkingStreamingEvent(ctx context.Context, payload *lifecycle.AgentStreamEventPayload) {
|
||||
if payload.Data.Text == "" || payload.SessionID == "" {
|
||||
return
|
||||
}
|
||||
|
||||
if s.messageCreator == nil {
|
||||
return
|
||||
}
|
||||
|
||||
// MessageID is pre-generated by the manager to avoid race conditions
|
||||
messageID := payload.Data.MessageID
|
||||
if messageID == "" {
|
||||
s.logger.Warn("streaming thinking event missing message ID",
|
||||
zap.String("task_id", payload.TaskID),
|
||||
zap.String("session_id", payload.SessionID))
|
||||
return
|
||||
}
|
||||
|
||||
if payload.Data.IsAppend {
|
||||
// Append to existing thinking message
|
||||
if err := s.messageCreator.AppendThinkingMessage(ctx, messageID, payload.Data.Text); err != nil {
|
||||
s.logger.Error("failed to append to streaming thinking message",
|
||||
zap.String("task_id", payload.TaskID),
|
||||
zap.String("message_id", messageID),
|
||||
zap.Error(err))
|
||||
} else {
|
||||
s.logger.Debug("appended to streaming thinking message",
|
||||
zap.String("task_id", payload.TaskID),
|
||||
zap.String("message_id", messageID),
|
||||
zap.Int("content_length", len(payload.Data.Text)))
|
||||
}
|
||||
} else {
|
||||
// Create new streaming thinking message with the pre-generated ID
|
||||
if err := s.messageCreator.CreateThinkingMessageStreaming(ctx, messageID, payload.TaskID, payload.Data.Text, payload.SessionID, s.getActiveTurnID(payload.SessionID)); err != nil {
|
||||
s.logger.Error("failed to create streaming thinking message",
|
||||
zap.String("task_id", payload.TaskID),
|
||||
zap.String("message_id", messageID),
|
||||
zap.Error(err))
|
||||
} else {
|
||||
s.logger.Debug("created streaming thinking message",
|
||||
zap.String("task_id", payload.TaskID),
|
||||
zap.String("session_id", payload.SessionID),
|
||||
zap.String("message_id", messageID),
|
||||
zap.Int("content_length", len(payload.Data.Text)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// handleToolUpdateEvent handles tool_update events and updates messages
|
||||
func (s *Service) handleToolUpdateEvent(ctx context.Context, payload *lifecycle.AgentStreamEventPayload) {
|
||||
if payload.SessionID == "" {
|
||||
|
||||
@@ -63,6 +63,10 @@ type MessageCreator interface {
|
||||
CreateAgentMessageStreaming(ctx context.Context, messageID, taskID, content, agentSessionID, turnID string) error
|
||||
// AppendAgentMessage appends additional content to an existing streaming message
|
||||
AppendAgentMessage(ctx context.Context, messageID, additionalContent string) error
|
||||
// CreateThinkingMessageStreaming creates a new thinking message with a pre-generated ID for streaming updates
|
||||
CreateThinkingMessageStreaming(ctx context.Context, messageID, taskID, content, agentSessionID, turnID string) error
|
||||
// AppendThinkingMessage appends additional content to an existing streaming thinking message
|
||||
AppendThinkingMessage(ctx context.Context, messageID, additionalContent string) error
|
||||
}
|
||||
|
||||
// TurnService is an interface for managing session turns
|
||||
|
||||
@@ -110,6 +110,8 @@ const (
|
||||
MessageTypePermissionRequest MessageType = "permission_request"
|
||||
// MessageTypeScriptExecution is for setup/cleanup script execution messages
|
||||
MessageTypeScriptExecution MessageType = "script_execution"
|
||||
// MessageTypeThinking is for agent thinking/reasoning content
|
||||
MessageTypeThinking MessageType = "thinking"
|
||||
)
|
||||
|
||||
// Message represents a message in a task session
|
||||
|
||||
@@ -1958,6 +1958,46 @@ func (s *Service) AppendMessageContent(ctx context.Context, messageID, additiona
|
||||
return nil
|
||||
}
|
||||
|
||||
// AppendThinkingContent appends additional thinking content to an existing thinking message.
|
||||
// This updates the metadata.thinking field for streaming agent reasoning.
|
||||
func (s *Service) AppendThinkingContent(ctx context.Context, messageID, additionalContent string) error {
|
||||
message, err := s.repo.GetMessage(ctx, messageID)
|
||||
if err != nil {
|
||||
s.logger.Warn("thinking message not found for append",
|
||||
zap.String("message_id", messageID),
|
||||
zap.Error(err))
|
||||
return err
|
||||
}
|
||||
|
||||
// Initialize metadata if nil
|
||||
if message.Metadata == nil {
|
||||
message.Metadata = make(map[string]interface{})
|
||||
}
|
||||
|
||||
// Get existing thinking content and append
|
||||
existingThinking := ""
|
||||
if existing, ok := message.Metadata["thinking"].(string); ok {
|
||||
existingThinking = existing
|
||||
}
|
||||
message.Metadata["thinking"] = existingThinking + additionalContent
|
||||
|
||||
if err := s.repo.UpdateMessage(ctx, message); err != nil {
|
||||
s.logger.Error("failed to append thinking content",
|
||||
zap.String("message_id", messageID),
|
||||
zap.Error(err))
|
||||
return err
|
||||
}
|
||||
|
||||
// Publish message.updated event for real-time streaming
|
||||
s.publishMessageEvent(ctx, events.MessageUpdated, message)
|
||||
|
||||
s.logger.Debug("thinking content appended",
|
||||
zap.String("message_id", messageID),
|
||||
zap.Int("appended_length", len(additionalContent)))
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateToolCallMessage updates a tool call message's status, and optionally title/args.
|
||||
// It includes retry logic to handle race conditions where the complete event
|
||||
// may arrive before the message has been created by the start event.
|
||||
|
||||
@@ -20,10 +20,6 @@ type Config struct {
|
||||
// Default: ~/.kandev/worktrees
|
||||
BasePath string `mapstructure:"base_path"`
|
||||
|
||||
// MaxPerRepo is the maximum number of concurrent worktrees per repository.
|
||||
// Default: 10
|
||||
MaxPerRepo int `mapstructure:"max_per_repo"`
|
||||
|
||||
// BranchPrefix is the prefix used for worktree branch names.
|
||||
// Default: feature/
|
||||
BranchPrefix string `mapstructure:"branch_prefix"`
|
||||
@@ -34,9 +30,6 @@ const DefaultBranchPrefix = "feature/"
|
||||
|
||||
// Validate validates the configuration and returns an error if invalid.
|
||||
func (c *Config) Validate() error {
|
||||
if c.MaxPerRepo <= 0 {
|
||||
c.MaxPerRepo = 10
|
||||
}
|
||||
if c.BranchPrefix == "" {
|
||||
c.BranchPrefix = DefaultBranchPrefix
|
||||
}
|
||||
|
||||
@@ -16,9 +16,6 @@ var (
|
||||
// ErrBranchExists is returned when the branch name already exists in the repository.
|
||||
ErrBranchExists = errors.New("branch already exists")
|
||||
|
||||
// ErrMaxWorktrees is returned when the maximum number of worktrees per repository is reached.
|
||||
ErrMaxWorktrees = errors.New("maximum worktrees reached for repository")
|
||||
|
||||
// ErrWorktreeLocked is returned when the worktree is locked by another process.
|
||||
ErrWorktreeLocked = errors.New("worktree is locked by another process")
|
||||
|
||||
|
||||
@@ -175,15 +175,6 @@ func (m *Manager) Create(ctx context.Context, req CreateRequest) (*Worktree, err
|
||||
return nil, ErrRepoNotGit
|
||||
}
|
||||
|
||||
// Check worktree limit
|
||||
count, err := m.countWorktreesForRepo(ctx, req.RepositoryID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to count worktrees: %w", err)
|
||||
}
|
||||
if count >= m.config.MaxPerRepo {
|
||||
return nil, fmt.Errorf("%w: %d", ErrMaxWorktrees, m.config.MaxPerRepo)
|
||||
}
|
||||
|
||||
// Get repository lock for safe concurrent access
|
||||
repoLock := m.getRepoLock(req.RepositoryPath)
|
||||
repoLock.Lock()
|
||||
@@ -716,35 +707,6 @@ func (m *Manager) pullBaseBranch(repoPath, baseBranch string) string {
|
||||
return baseBranch
|
||||
}
|
||||
|
||||
// countWorktreesForRepo counts active worktrees for a repository.
|
||||
func (m *Manager) countWorktreesForRepo(ctx context.Context, repoID string) (int, error) {
|
||||
if m.store == nil {
|
||||
// Count from cache if no store
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
count := 0
|
||||
for _, wt := range m.worktrees {
|
||||
if wt.RepositoryID == repoID && wt.Status == StatusActive {
|
||||
count++
|
||||
}
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
worktrees, err := m.store.GetWorktreesByRepositoryID(ctx, repoID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
count := 0
|
||||
for _, wt := range worktrees {
|
||||
if wt.Status == StatusActive {
|
||||
count++
|
||||
}
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
// removeWorktreeDir removes a worktree directory using git worktree remove.
|
||||
func (m *Manager) removeWorktreeDir(ctx context.Context, worktreePath, repoPath string) error {
|
||||
// First try git worktree remove
|
||||
|
||||
@@ -23,7 +23,6 @@ func newTestConfig(t *testing.T) Config {
|
||||
return Config{
|
||||
Enabled: true,
|
||||
BasePath: tmpDir,
|
||||
MaxPerRepo: 10,
|
||||
BranchPrefix: "kandev/",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -16,7 +16,6 @@ func Provide(db *sql.DB, cfg *config.Config, log *logger.Logger) (*Manager, func
|
||||
manager, err := NewManager(Config{
|
||||
Enabled: cfg.Worktree.Enabled,
|
||||
BasePath: cfg.Worktree.BasePath,
|
||||
MaxPerRepo: cfg.Worktree.MaxPerRepo,
|
||||
BranchPrefix: "kandev/",
|
||||
}, store, log)
|
||||
if err != nil {
|
||||
|
||||
Reference in New Issue
Block a user