mirror of
https://github.com/mark3labs/kit.git
synced 2026-06-17 13:06:32 +00:00
Compare commits
4 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 558fb5214f | |||
| 61408ed490 | |||
| 3cfb6437f9 | |||
| d33ad4028b |
@@ -29,7 +29,7 @@ A powerful, extensible AI coding agent CLI with multi-provider support, built-in
|
||||
- **Session Management**: Tree-based conversation history with branching support
|
||||
- **Non-Interactive Mode**: Script-friendly positional args with JSON output
|
||||
- **ACP Server**: Run Kit as an [Agent Client Protocol](https://agentclientprotocol.com) agent over stdio
|
||||
- **Go SDK**: Embed Kit in your own applications
|
||||
- **Go SDK**: Embed Kit in your own applications with full agent lifecycle events (30+ event types) and behavior-modifying hooks
|
||||
|
||||
## Installation
|
||||
|
||||
@@ -646,7 +646,28 @@ host, _ := kit.New(ctx, &kit.Options{
|
||||
})
|
||||
```
|
||||
|
||||
Use `kit.NewParallelTool` for tools safe to run concurrently. See the [SDK docs](/sdk/overview) for full details on struct tags, `ToolOutput` fields, and `ToolCallIDFromContext`.
|
||||
Use `kit.NewParallelTool` for tools safe to run concurrently. Binary data (images, audio, etc.) in `ToolOutput.Data` is automatically forwarded to the LLM when `MediaType` is set. See the [SDK docs](/sdk/overview) for full details on struct tags, `ToolOutput` fields, and `ToolCallIDFromContext`.
|
||||
|
||||
#### Return Helpers
|
||||
|
||||
| Helper | Description |
|
||||
| --- | --- |
|
||||
| `kit.TextResult(content)` | Successful text result |
|
||||
| `kit.ErrorResult(content)` | Error result (LLM sees it as a tool error) |
|
||||
| `kit.ImageResult(content, data, mediaType)` | Image result with binary data (e.g. `"image/png"`) |
|
||||
| `kit.MediaResult(content, data, mediaType)` | Non-image media result (e.g. `"audio/mpeg"`) |
|
||||
|
||||
#### ToolOutput Fields
|
||||
|
||||
```go
|
||||
kit.ToolOutput{
|
||||
Content: "result text", // text returned to the LLM
|
||||
IsError: false, // true = LLM sees this as an error
|
||||
Data: pngBytes, // optional binary data (images, audio)
|
||||
MediaType: "image/png", // MIME type for binary Data
|
||||
Metadata: map[string]any{}, // opaque metadata for hooks/UI (not sent to LLM)
|
||||
}
|
||||
```
|
||||
|
||||
### With Callbacks
|
||||
|
||||
@@ -663,7 +684,7 @@ unsub2 := host.OnToolResult(func(e kit.ToolResultEvent) {
|
||||
})
|
||||
defer unsub2()
|
||||
|
||||
unsub3 := host.OnStreaming(func(e kit.MessageUpdateEvent) {
|
||||
unsub3 := host.OnMessageUpdate(func(e kit.MessageUpdateEvent) {
|
||||
print(e.Chunk)
|
||||
})
|
||||
defer unsub3()
|
||||
|
||||
@@ -62,7 +62,7 @@ func main() {
|
||||
}
|
||||
})
|
||||
// Subscribe to streaming chunks.
|
||||
host3.OnStreaming(func(e kit.MessageUpdateEvent) {
|
||||
host3.OnMessageUpdate(func(e kit.MessageUpdateEvent) {
|
||||
fmt.Print(e.Chunk)
|
||||
})
|
||||
|
||||
|
||||
+289
-70
@@ -6,6 +6,7 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"charm.land/fantasy"
|
||||
|
||||
@@ -126,6 +127,76 @@ type StepMessagesHandler func(stepMessages []fantasy.Message)
|
||||
// tracking during long-running tool-calling conversations.
|
||||
type StepUsageHandler func(inputTokens, outputTokens, cacheReadTokens, cacheCreationTokens int64)
|
||||
|
||||
// StepStartHandler is called when a new LLM step begins within a turn.
|
||||
type StepStartHandler func(stepNumber int)
|
||||
|
||||
// StepFinishHandler is called when a step completes with full context.
|
||||
type StepFinishHandler func(stepNumber int, hasToolCalls bool, finishReason string, usage fantasy.Usage)
|
||||
|
||||
// TextStartHandler is called when the LLM begins generating text content.
|
||||
type TextStartHandler func(id string)
|
||||
|
||||
// TextEndHandler is called when the LLM finishes generating text content.
|
||||
type TextEndHandler func(id string)
|
||||
|
||||
// ReasoningStartHandler is called when the LLM begins reasoning/thinking.
|
||||
type ReasoningStartHandler func(id string)
|
||||
|
||||
// WarningsHandler is called when the LLM provider returns warnings.
|
||||
type WarningsHandler func(warnings []string)
|
||||
|
||||
// SourceHandler is called when the LLM references a source.
|
||||
type SourceHandler func(sourceType, id, url, title string)
|
||||
|
||||
// StreamFinishHandler is called when a per-step LLM stream completes.
|
||||
type StreamFinishHandler func(usage fantasy.Usage, finishReason string)
|
||||
|
||||
// ErrorHandler is called when an agent-level error occurs.
|
||||
type ErrorHandler func(err error)
|
||||
|
||||
// RetryHandler is called when the LLM request is retried.
|
||||
type RetryHandler func(attempt int, err error)
|
||||
|
||||
// PrepareStepHandler is called between steps to allow message modification.
|
||||
// It receives the step number and current messages, and returns replacement
|
||||
// messages (or nil to keep unchanged).
|
||||
type PrepareStepHandler func(stepNumber int, messages []fantasy.Message) []fantasy.Message
|
||||
|
||||
// GenerateCallbacks consolidates all callback functions for
|
||||
// GenerateWithLoopAndStreaming into a single struct. This replaces the previous
|
||||
// 16+ positional callback parameters, making it easier to add new callbacks
|
||||
// without breaking existing callers (new fields default to nil).
|
||||
type GenerateCallbacks struct {
|
||||
OnToolCall ToolCallHandler
|
||||
OnToolExecution ToolExecutionHandler
|
||||
OnToolResult ToolResultHandler
|
||||
OnResponse ResponseHandler
|
||||
OnToolCallContent ToolCallContentHandler
|
||||
OnStreamingResponse StreamingResponseHandler
|
||||
OnReasoningDelta ReasoningDeltaHandler
|
||||
OnReasoningComplete ReasoningCompleteHandler
|
||||
OnToolOutput ToolOutputHandler
|
||||
OnStepMessages StepMessagesHandler
|
||||
OnStepUsage StepUsageHandler
|
||||
OnPasswordPrompt PasswordPromptHandler
|
||||
OnToolCallStart ToolCallStartHandler
|
||||
OnToolCallDelta ToolCallDeltaHandler
|
||||
OnToolCallEnd ToolCallEndHandler
|
||||
|
||||
// New callbacks for previously unwired Fantasy lifecycle events.
|
||||
OnStepStart StepStartHandler
|
||||
OnStepFinish StepFinishHandler
|
||||
OnTextStart TextStartHandler
|
||||
OnTextEnd TextEndHandler
|
||||
OnReasoningStart ReasoningStartHandler
|
||||
OnWarnings WarningsHandler
|
||||
OnSource SourceHandler
|
||||
OnStreamFinish StreamFinishHandler
|
||||
OnError ErrorHandler
|
||||
OnRetry RetryHandler
|
||||
OnPrepareStep PrepareStepHandler
|
||||
}
|
||||
|
||||
// Agent represents an AI agent with core tool integration using the LLM library.
|
||||
// Core tools (bash, read, write, edit, grep, find, ls) are registered as direct
|
||||
// AgentTool implementations — no MCP layer, no serialization overhead.
|
||||
@@ -423,13 +494,20 @@ func (a *Agent) GenerateWithLoop(ctx context.Context, messages []fantasy.Message
|
||||
onToolCall ToolCallHandler, onToolExecution ToolExecutionHandler, onToolResult ToolResultHandler,
|
||||
onResponse ResponseHandler, onToolCallContent ToolCallContentHandler,
|
||||
) (*GenerateWithLoopResult, error) {
|
||||
return a.GenerateWithLoopAndStreaming(ctx, messages, onToolCall, onToolExecution, onToolResult,
|
||||
onResponse, onToolCallContent, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
return a.GenerateWithCallbacks(ctx, messages, GenerateCallbacks{
|
||||
OnToolCall: onToolCall,
|
||||
OnToolExecution: onToolExecution,
|
||||
OnToolResult: onToolResult,
|
||||
OnResponse: onResponse,
|
||||
OnToolCallContent: onToolCallContent,
|
||||
})
|
||||
}
|
||||
|
||||
// GenerateWithLoopAndStreaming processes messages using the agent with streaming and callbacks.
|
||||
// The agent handles the tool call loop internally. We map the rich callback system
|
||||
// to kit's existing callback interface for UI integration.
|
||||
// The agent handles the tool call loop internally.
|
||||
//
|
||||
// Deprecated: Use GenerateWithCallbacks instead, which takes a GenerateCallbacks
|
||||
// struct and is easier to extend with new callbacks.
|
||||
func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fantasy.Message,
|
||||
onToolCall ToolCallHandler, onToolExecution ToolExecutionHandler, onToolResult ToolResultHandler,
|
||||
onResponse ResponseHandler, onToolCallContent ToolCallContentHandler,
|
||||
@@ -444,6 +522,31 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
|
||||
onToolCallDelta ToolCallDeltaHandler,
|
||||
onToolCallEnd ToolCallEndHandler,
|
||||
) (*GenerateWithLoopResult, error) {
|
||||
return a.GenerateWithCallbacks(ctx, messages, GenerateCallbacks{
|
||||
OnToolCall: onToolCall,
|
||||
OnToolExecution: onToolExecution,
|
||||
OnToolResult: onToolResult,
|
||||
OnResponse: onResponse,
|
||||
OnToolCallContent: onToolCallContent,
|
||||
OnStreamingResponse: onStreamingResponse,
|
||||
OnReasoningDelta: onReasoningDelta,
|
||||
OnReasoningComplete: onReasoningComplete,
|
||||
OnToolOutput: onToolOutput,
|
||||
OnStepMessages: onStepMessages,
|
||||
OnStepUsage: onStepUsage,
|
||||
OnPasswordPrompt: onPasswordPrompt,
|
||||
OnToolCallStart: onToolCallStart,
|
||||
OnToolCallDelta: onToolCallDelta,
|
||||
OnToolCallEnd: onToolCallEnd,
|
||||
})
|
||||
}
|
||||
|
||||
// GenerateWithCallbacks processes messages using the agent with streaming and callbacks.
|
||||
// The agent handles the tool call loop internally. We map the rich callback system
|
||||
// to kit's existing callback interface for UI integration.
|
||||
func (a *Agent) GenerateWithCallbacks(ctx context.Context, messages []fantasy.Message,
|
||||
cb GenerateCallbacks,
|
||||
) (*GenerateWithLoopResult, error) {
|
||||
|
||||
// Wait for background MCP tool loading to complete and rebuild the
|
||||
// fantasy agent with the full tool set. This is a no-op when no MCP
|
||||
@@ -451,13 +554,13 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
|
||||
a.ensureMCPTools()
|
||||
|
||||
// Inject tool output handler into context for use by core tools (e.g., bash).
|
||||
if onToolOutput != nil {
|
||||
ctx = core.ContextWithToolOutputCallback(ctx, onToolOutput)
|
||||
if cb.OnToolOutput != nil {
|
||||
ctx = core.ContextWithToolOutputCallback(ctx, cb.OnToolOutput)
|
||||
}
|
||||
|
||||
// Inject password prompt handler into context for use by bash tool.
|
||||
if onPasswordPrompt != nil {
|
||||
ctx = core.ContextWithPasswordPrompt(ctx, onPasswordPrompt)
|
||||
if cb.OnPasswordPrompt != nil {
|
||||
ctx = core.ContextWithPasswordPrompt(ctx, cb.OnPasswordPrompt)
|
||||
}
|
||||
|
||||
// The agent requires the current user input as Prompt, with prior messages as history.
|
||||
@@ -477,9 +580,13 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
|
||||
// provided. The agent only exposes tool/step callbacks on AgentStreamCall, so
|
||||
// Stream is required to observe tool execution in real time. The non-streaming
|
||||
// Generate path is reserved for the simple case with no callbacks at all.
|
||||
hasCallbacks := onToolCall != nil || onToolExecution != nil || onToolResult != nil ||
|
||||
onToolCallContent != nil || onStreamingResponse != nil || onReasoningDelta != nil ||
|
||||
onToolCallStart != nil || onToolCallDelta != nil || onToolCallEnd != nil
|
||||
hasCallbacks := cb.OnToolCall != nil || cb.OnToolExecution != nil || cb.OnToolResult != nil ||
|
||||
cb.OnToolCallContent != nil || cb.OnStreamingResponse != nil || cb.OnReasoningDelta != nil ||
|
||||
cb.OnToolCallStart != nil || cb.OnToolCallDelta != nil || cb.OnToolCallEnd != nil ||
|
||||
cb.OnStepStart != nil || cb.OnStepFinish != nil || cb.OnTextStart != nil ||
|
||||
cb.OnTextEnd != nil || cb.OnReasoningStart != nil || cb.OnWarnings != nil ||
|
||||
cb.OnSource != nil || cb.OnStreamFinish != nil || cb.OnError != nil ||
|
||||
cb.OnRetry != nil || cb.OnPrepareStep != nil
|
||||
|
||||
if a.streamingEnabled || hasCallbacks {
|
||||
// Track completed step messages so we can return partial results
|
||||
@@ -488,9 +595,11 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
|
||||
// for every step that completed before the error occurred.
|
||||
var completedStepMessages []fantasy.Message
|
||||
// persistedCount tracks how many new messages (beyond the original
|
||||
// input) were persisted incrementally via onStepMessages, so the
|
||||
// input) were persisted incrementally via cb.OnStepMessages, so the
|
||||
// caller can skip them during post-generation persistence.
|
||||
var persistedCount int
|
||||
// stepCounter tracks the current step number for StepStart/StepFinish events.
|
||||
var stepCounter int
|
||||
|
||||
// Use the streaming agent
|
||||
streamCall := fantasy.AgentStreamCall{
|
||||
@@ -503,8 +612,8 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
if onToolCallStart != nil {
|
||||
onToolCallStart(id, toolName)
|
||||
if cb.OnToolCallStart != nil {
|
||||
cb.OnToolCallStart(id, toolName)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
@@ -512,8 +621,8 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
if onToolCallDelta != nil {
|
||||
onToolCallDelta(id, delta)
|
||||
if cb.OnToolCallDelta != nil {
|
||||
cb.OnToolCallDelta(id, delta)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
@@ -521,8 +630,39 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
if onToolCallEnd != nil {
|
||||
onToolCallEnd(id)
|
||||
if cb.OnToolCallEnd != nil {
|
||||
cb.OnToolCallEnd(id)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
|
||||
// Text start/end callbacks
|
||||
OnTextStart: func(id string) error {
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
if cb.OnTextStart != nil {
|
||||
cb.OnTextStart(id)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
OnTextEnd: func(id string) error {
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
if cb.OnTextEnd != nil {
|
||||
cb.OnTextEnd(id)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
|
||||
// Reasoning start callback
|
||||
OnReasoningStart: func(id string, _ fantasy.ReasoningContent) error {
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
if cb.OnReasoningStart != nil {
|
||||
cb.OnReasoningStart(id)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
@@ -532,8 +672,8 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
if onReasoningDelta != nil {
|
||||
onReasoningDelta(delta)
|
||||
if cb.OnReasoningDelta != nil {
|
||||
cb.OnReasoningDelta(delta)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
@@ -543,8 +683,8 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
if onReasoningComplete != nil {
|
||||
onReasoningComplete()
|
||||
if cb.OnReasoningComplete != nil {
|
||||
cb.OnReasoningComplete()
|
||||
}
|
||||
return nil
|
||||
},
|
||||
@@ -554,8 +694,64 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
if onStreamingResponse != nil {
|
||||
onStreamingResponse(text)
|
||||
if cb.OnStreamingResponse != nil {
|
||||
cb.OnStreamingResponse(text)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
|
||||
// Warnings callback
|
||||
OnWarnings: func(warnings []fantasy.CallWarning) error {
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
if cb.OnWarnings != nil {
|
||||
strs := make([]string, len(warnings))
|
||||
for i, w := range warnings {
|
||||
strs[i] = w.Message
|
||||
}
|
||||
cb.OnWarnings(strs)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
|
||||
// Source callback
|
||||
OnSource: func(source fantasy.SourceContent) error {
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
if cb.OnSource != nil {
|
||||
cb.OnSource(string(source.SourceType), source.ID, source.URL, source.Title)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
|
||||
// Stream finish callback (per-step stream completion)
|
||||
OnStreamFinish: func(usage fantasy.Usage, finishReason fantasy.FinishReason, _ fantasy.ProviderMetadata) error {
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
if cb.OnStreamFinish != nil {
|
||||
cb.OnStreamFinish(usage, string(finishReason))
|
||||
}
|
||||
return nil
|
||||
},
|
||||
|
||||
// Error callback
|
||||
OnError: func(err error) {
|
||||
if cb.OnError != nil {
|
||||
cb.OnError(err)
|
||||
}
|
||||
},
|
||||
|
||||
// Step start callback
|
||||
OnStepStart: func(stepNumber int) error {
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
stepCounter = stepNumber
|
||||
if cb.OnStepStart != nil {
|
||||
cb.OnStepStart(stepNumber)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
@@ -568,13 +764,13 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
|
||||
currentToolArgs = tc.Input
|
||||
|
||||
// Notify about the tool call
|
||||
if onToolCall != nil {
|
||||
onToolCall(tc.ToolCallID, tc.ToolName, tc.Input)
|
||||
if cb.OnToolCall != nil {
|
||||
cb.OnToolCall(tc.ToolCallID, tc.ToolName, tc.Input)
|
||||
}
|
||||
|
||||
// Notify tool execution starting
|
||||
if onToolExecution != nil {
|
||||
onToolExecution(tc.ToolCallID, tc.ToolName, tc.Input, true)
|
||||
if cb.OnToolExecution != nil {
|
||||
cb.OnToolExecution(tc.ToolCallID, tc.ToolName, tc.Input, true)
|
||||
}
|
||||
|
||||
return nil
|
||||
@@ -586,14 +782,14 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
|
||||
return ctx.Err()
|
||||
}
|
||||
// Notify tool execution finished
|
||||
if onToolExecution != nil {
|
||||
onToolExecution(tr.ToolCallID, tr.ToolName, currentToolArgs, false)
|
||||
if cb.OnToolExecution != nil {
|
||||
cb.OnToolExecution(tr.ToolCallID, tr.ToolName, currentToolArgs, false)
|
||||
}
|
||||
|
||||
if onToolResult != nil {
|
||||
if cb.OnToolResult != nil {
|
||||
// Extract result text and error status
|
||||
resultText, isError := extractToolResultText(tr)
|
||||
onToolResult(tr.ToolCallID, tr.ToolName, currentToolArgs, resultText, tr.ClientMetadata, isError)
|
||||
cb.OnToolResult(tr.ToolCallID, tr.ToolName, currentToolArgs, resultText, tr.ClientMetadata, isError)
|
||||
}
|
||||
|
||||
return nil
|
||||
@@ -607,8 +803,8 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
|
||||
|
||||
// Persist step messages incrementally so progress is saved
|
||||
// as it happens rather than only at the end of the turn.
|
||||
if onStepMessages != nil && len(step.Messages) > 0 {
|
||||
onStepMessages(step.Messages)
|
||||
if cb.OnStepMessages != nil && len(step.Messages) > 0 {
|
||||
cb.OnStepMessages(step.Messages)
|
||||
persistedCount += len(step.Messages)
|
||||
}
|
||||
|
||||
@@ -618,65 +814,88 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
|
||||
// Check if step has text content alongside tool calls
|
||||
text := step.Content.Text()
|
||||
toolCalls := step.Content.ToolCalls()
|
||||
if text != "" && len(toolCalls) > 0 && onToolCallContent != nil {
|
||||
onToolCallContent(text)
|
||||
if text != "" && len(toolCalls) > 0 && cb.OnToolCallContent != nil {
|
||||
cb.OnToolCallContent(text)
|
||||
}
|
||||
// Emit step usage for real-time cost tracking
|
||||
if onStepUsage != nil {
|
||||
onStepUsage(step.Usage.InputTokens, step.Usage.OutputTokens,
|
||||
if cb.OnStepUsage != nil {
|
||||
cb.OnStepUsage(step.Usage.InputTokens, step.Usage.OutputTokens,
|
||||
step.Usage.CacheReadTokens, step.Usage.CacheCreationTokens)
|
||||
}
|
||||
// Emit unified step finish event
|
||||
if cb.OnStepFinish != nil {
|
||||
cb.OnStepFinish(stepCounter, len(toolCalls) > 0, string(step.FinishReason), step.Usage)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
||||
// If a steer channel is attached to the context, wire up a
|
||||
// PrepareStep function that drains the channel between steps
|
||||
// and injects pending steer messages as user messages before
|
||||
// the next LLM call. This enables graceful mid-turn steering
|
||||
// without cancelling in-progress tool execution.
|
||||
if steerCh := steerChFromContext(ctx); steerCh != nil {
|
||||
onConsumed := steerConsumedFromContext(ctx)
|
||||
// Always wire up PrepareStep to handle both steering and the
|
||||
// OnPrepareStep hook. Steering drains its channel first, then
|
||||
// OnPrepareStep hooks run against the (possibly already steered)
|
||||
// messages.
|
||||
steerCh := steerChFromContext(ctx)
|
||||
onConsumed := steerConsumedFromContext(ctx)
|
||||
hasSteering := steerCh != nil
|
||||
hasPrepareStepHook := cb.OnPrepareStep != nil
|
||||
|
||||
if hasSteering || hasPrepareStepHook {
|
||||
streamCall.PrepareStep = func(
|
||||
stepCtx context.Context,
|
||||
opts fantasy.PrepareStepFunctionOptions,
|
||||
) (context.Context, fantasy.PrepareStepResult, error) {
|
||||
// Drain all pending steer messages (non-blocking).
|
||||
var steered []SteerMessage
|
||||
for {
|
||||
select {
|
||||
case msg := <-steerCh:
|
||||
steered = append(steered, msg)
|
||||
default:
|
||||
goto done
|
||||
}
|
||||
}
|
||||
done:
|
||||
result := fantasy.PrepareStepResult{
|
||||
Model: opts.Model,
|
||||
Messages: opts.Messages,
|
||||
}
|
||||
if len(steered) > 0 {
|
||||
// Inject each steer message as a user message so the
|
||||
// LLM sees the redirection on the next step.
|
||||
for _, sm := range steered {
|
||||
result.Messages = append(result.Messages,
|
||||
fantasy.NewUserMessage(sm.Text, sm.Files...))
|
||||
|
||||
// Phase 1: Drain steering channel (if present).
|
||||
if hasSteering {
|
||||
var steered []SteerMessage
|
||||
for {
|
||||
select {
|
||||
case msg := <-steerCh:
|
||||
steered = append(steered, msg)
|
||||
default:
|
||||
goto done
|
||||
}
|
||||
}
|
||||
// Notify that steer messages were consumed.
|
||||
if onConsumed != nil {
|
||||
onConsumed(len(steered))
|
||||
done:
|
||||
if len(steered) > 0 {
|
||||
for _, sm := range steered {
|
||||
result.Messages = append(result.Messages,
|
||||
fantasy.NewUserMessage(sm.Text, sm.Files...))
|
||||
}
|
||||
if onConsumed != nil {
|
||||
onConsumed(len(steered))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 2: Run OnPrepareStep hook (if registered).
|
||||
if hasPrepareStepHook {
|
||||
if replacement := cb.OnPrepareStep(opts.StepNumber, result.Messages); replacement != nil {
|
||||
result.Messages = replacement
|
||||
}
|
||||
}
|
||||
|
||||
// Apply message-level cache control for Anthropic models.
|
||||
// This avoids type conflicts with provider-level options.
|
||||
result.Messages = applyCacheControlToMessages(result.Messages)
|
||||
|
||||
return stepCtx, result, nil
|
||||
}
|
||||
}
|
||||
|
||||
// Wire OnRetry callback if provided.
|
||||
if cb.OnRetry != nil {
|
||||
streamCall.OnRetry = func(err *fantasy.ProviderError, _ time.Duration) {
|
||||
// Use the retry number from the error if available; Fantasy
|
||||
// doesn't pass a counter directly, so we approximate with a
|
||||
// counter incremented on each call.
|
||||
cb.OnRetry(0, err)
|
||||
}
|
||||
}
|
||||
|
||||
result, err := a.fantasyAgent.Stream(ctx, streamCall)
|
||||
if err != nil {
|
||||
// On cancellation (or any error), return a partial result
|
||||
@@ -702,8 +921,8 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
|
||||
// empty (e.g. reasoning-only responses) so the UI properly resets
|
||||
// the stream component and avoids duplicate content on the next
|
||||
// flush.
|
||||
if onResponse != nil {
|
||||
onResponse(result.Response.Content.Text())
|
||||
if cb.OnResponse != nil {
|
||||
cb.OnResponse(result.Response.Content.Text())
|
||||
}
|
||||
|
||||
r := convertAgentResult(result, messages)
|
||||
@@ -723,8 +942,8 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
|
||||
|
||||
// For non-streaming, fire the response callback so callers can reset
|
||||
// streaming state (see streaming path comment above).
|
||||
if onResponse != nil {
|
||||
onResponse(result.Response.Content.Text())
|
||||
if cb.OnResponse != nil {
|
||||
cb.OnResponse(result.Response.Content.Text())
|
||||
}
|
||||
|
||||
return convertAgentResult(result, messages), nil
|
||||
|
||||
@@ -1094,6 +1094,14 @@ type API struct {
|
||||
onSubagentStart func(func(SubagentStartEvent, Context))
|
||||
onSubagentChunk func(func(SubagentChunkEvent, Context))
|
||||
onSubagentEnd func(func(SubagentEndEvent, Context))
|
||||
onStepStart func(func(StepStartEvent, Context))
|
||||
onStepFinish func(func(StepFinishEvent, Context))
|
||||
onReasoningStart func(func(ReasoningStartEvent, Context))
|
||||
onWarnings func(func(WarningsEvent, Context))
|
||||
onSource func(func(SourceEvent, Context))
|
||||
onError func(func(ErrorEvent, Context))
|
||||
onRetry func(func(RetryEvent, Context))
|
||||
onPrepareStep func(func(PrepareStepEvent, Context) *PrepareStepResult)
|
||||
}
|
||||
|
||||
// OnToolCall registers a handler that fires before a tool executes.
|
||||
@@ -1301,6 +1309,56 @@ func (a *API) OnBeforeCompact(handler func(BeforeCompactEvent, Context) *BeforeC
|
||||
a.onBeforeCompact(handler)
|
||||
}
|
||||
|
||||
// OnStepStart registers a handler that fires when a new LLM call begins
|
||||
// within a multi-step agent turn.
|
||||
func (a *API) OnStepStart(handler func(StepStartEvent, Context)) {
|
||||
a.onStepStart(handler)
|
||||
}
|
||||
|
||||
// OnStepFinish registers a handler that fires when a step completes,
|
||||
// providing step number, finish reason, and decomposed token usage.
|
||||
func (a *API) OnStepFinish(handler func(StepFinishEvent, Context)) {
|
||||
a.onStepFinish(handler)
|
||||
}
|
||||
|
||||
// OnReasoningStart registers a handler that fires when the LLM begins
|
||||
// reasoning/thinking.
|
||||
func (a *API) OnReasoningStart(handler func(ReasoningStartEvent, Context)) {
|
||||
a.onReasoningStart(handler)
|
||||
}
|
||||
|
||||
// OnWarnings registers a handler that fires when the LLM provider returns
|
||||
// warnings about the request.
|
||||
func (a *API) OnWarnings(handler func(WarningsEvent, Context)) {
|
||||
a.onWarnings(handler)
|
||||
}
|
||||
|
||||
// OnSource registers a handler that fires when the LLM references a source
|
||||
// (e.g. from web search tools).
|
||||
func (a *API) OnSource(handler func(SourceEvent, Context)) {
|
||||
a.onSource(handler)
|
||||
}
|
||||
|
||||
// OnError registers a handler that fires when an agent-level error occurs
|
||||
// during streaming.
|
||||
func (a *API) OnError(handler func(ErrorEvent, Context)) {
|
||||
a.onError(handler)
|
||||
}
|
||||
|
||||
// OnRetry registers a handler that fires when the LLM provider request is
|
||||
// retried after a transient error.
|
||||
func (a *API) OnRetry(handler func(RetryEvent, Context)) {
|
||||
a.onRetry(handler)
|
||||
}
|
||||
|
||||
// OnPrepareStep registers a handler that fires between steps within a
|
||||
// multi-step agent turn, after steering messages are injected and before
|
||||
// messages are sent to the LLM. Return a non-nil PrepareStepResult with
|
||||
// Messages to replace the context window for this step.
|
||||
func (a *API) OnPrepareStep(handler func(PrepareStepEvent, Context) *PrepareStepResult) {
|
||||
a.onPrepareStep(handler)
|
||||
}
|
||||
|
||||
// RegisterToolRenderer registers a custom renderer for a specific tool's
|
||||
// display in the TUI. The renderer controls the header (parameter summary)
|
||||
// and/or body (result display) of the tool's output block. If multiple
|
||||
@@ -2253,6 +2311,98 @@ type SubagentEndEvent struct {
|
||||
|
||||
func (e SubagentEndEvent) Type() EventType { return SubagentEnd }
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Step lifecycle events (exposed to Yaegi — concrete structs)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// StepStartEvent fires when a new LLM call begins within a multi-step agent turn.
|
||||
type StepStartEvent struct {
|
||||
StepNumber int
|
||||
}
|
||||
|
||||
func (e StepStartEvent) Type() EventType { return StepStart }
|
||||
|
||||
// StepFinishEvent fires when a step completes, providing step metadata and
|
||||
// token usage. Usage fields are plain int64 (not LLMUsage) because Yaegi
|
||||
// cannot handle fantasy types across the interpreter boundary.
|
||||
type StepFinishEvent struct {
|
||||
StepNumber int
|
||||
HasToolCalls bool
|
||||
FinishReason string
|
||||
InputTokens int64
|
||||
OutputTokens int64
|
||||
CacheReadTokens int64
|
||||
CacheWriteTokens int64
|
||||
}
|
||||
|
||||
func (e StepFinishEvent) Type() EventType { return StepFinish }
|
||||
|
||||
// ReasoningStartEvent fires when the LLM begins reasoning/thinking.
|
||||
type ReasoningStartEvent struct {
|
||||
ID string
|
||||
}
|
||||
|
||||
func (e ReasoningStartEvent) Type() EventType { return ReasoningStart }
|
||||
|
||||
// WarningsEvent fires when the LLM provider returns warnings about the request.
|
||||
type WarningsEvent struct {
|
||||
Warnings []string
|
||||
}
|
||||
|
||||
func (e WarningsEvent) Type() EventType { return Warnings }
|
||||
|
||||
// SourceEvent fires when the LLM references a source (e.g. from web search).
|
||||
type SourceEvent struct {
|
||||
SourceType string
|
||||
ID string
|
||||
URL string
|
||||
Title string
|
||||
}
|
||||
|
||||
func (e SourceEvent) Type() EventType { return Source }
|
||||
|
||||
// ErrorEvent fires when an agent-level error occurs during streaming.
|
||||
// Uses string instead of error because Yaegi cannot handle the error
|
||||
// interface reliably across the interpreter boundary.
|
||||
type ErrorEvent struct {
|
||||
Error string
|
||||
}
|
||||
|
||||
func (e ErrorEvent) Type() EventType { return Error }
|
||||
|
||||
// RetryEvent fires when the LLM provider request is retried after a
|
||||
// transient error.
|
||||
type RetryEvent struct {
|
||||
Attempt int
|
||||
Error string
|
||||
}
|
||||
|
||||
func (e RetryEvent) Type() EventType { return Retry }
|
||||
|
||||
// PrepareStepEvent fires between steps within a multi-step agent turn,
|
||||
// after steering messages are injected and before messages are sent to
|
||||
// the LLM. Handlers can inspect and replace the context window.
|
||||
type PrepareStepEvent struct {
|
||||
// StepNumber is the zero-based step index within the current turn.
|
||||
StepNumber int
|
||||
// Messages is the current context window that will be sent to the LLM.
|
||||
Messages []ContextMessage
|
||||
}
|
||||
|
||||
func (e PrepareStepEvent) Type() EventType { return PrepareStep }
|
||||
|
||||
// PrepareStepResult allows extensions to replace the context window between
|
||||
// steps. Return nil Messages to leave the context unchanged.
|
||||
type PrepareStepResult struct {
|
||||
// Messages replaces the entire context window for this step. If nil,
|
||||
// the original messages are used unchanged. Messages with a non-negative
|
||||
// Index reuse the original message at that position; messages with
|
||||
// Index < 0 are created fresh from Role + Content.
|
||||
Messages []ContextMessage
|
||||
}
|
||||
|
||||
func (PrepareStepResult) isResult() {}
|
||||
|
||||
// ThemeColor is an adaptive color pair with light and dark hex values.
|
||||
// Either field may be empty to inherit from the default theme.
|
||||
type ThemeColor struct {
|
||||
|
||||
@@ -96,6 +96,35 @@ const (
|
||||
// SubagentEnd fires when a subagent tool call completes (success
|
||||
// or error). Carries the final response and any error message.
|
||||
SubagentEnd EventType = "subagent_end"
|
||||
|
||||
// StepStart fires when a new LLM call begins within a multi-step
|
||||
// agent turn.
|
||||
StepStart EventType = "step_start"
|
||||
|
||||
// StepFinish fires when a step completes, providing step number,
|
||||
// finish reason, and token usage.
|
||||
StepFinish EventType = "step_finish"
|
||||
|
||||
// ReasoningStart fires when the LLM begins reasoning/thinking.
|
||||
ReasoningStart EventType = "reasoning_start"
|
||||
|
||||
// Warnings fires when the LLM provider returns warnings.
|
||||
Warnings EventType = "warnings"
|
||||
|
||||
// Source fires when the LLM references a source (e.g. web search).
|
||||
Source EventType = "source"
|
||||
|
||||
// Error fires when an agent-level error occurs during streaming.
|
||||
Error EventType = "error"
|
||||
|
||||
// Retry fires when the LLM provider request is retried after a
|
||||
// transient error.
|
||||
Retry EventType = "retry"
|
||||
|
||||
// PrepareStep fires between steps within a multi-step agent turn,
|
||||
// after steering messages are injected and before messages are sent
|
||||
// to the LLM. Handlers can replace the context window for this step.
|
||||
PrepareStep EventType = "prepare_step"
|
||||
)
|
||||
|
||||
// AllEventTypes returns every supported event type.
|
||||
@@ -109,6 +138,8 @@ func AllEventTypes() []EventType {
|
||||
ModelChange, ContextPrepare,
|
||||
BeforeFork, BeforeSessionSwitch, BeforeCompact,
|
||||
SubagentStart, SubagentChunk, SubagentEnd,
|
||||
StepStart, StepFinish, ReasoningStart, Warnings, Source, Error, Retry,
|
||||
PrepareStep,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -4,8 +4,8 @@ import "testing"
|
||||
|
||||
func TestAllEventTypes_Count(t *testing.T) {
|
||||
all := AllEventTypes()
|
||||
if len(all) != 24 {
|
||||
t.Fatalf("expected 24 event types, got %d", len(all))
|
||||
if len(all) != 32 {
|
||||
t.Fatalf("expected 32 event types, got %d", len(all))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -618,6 +618,57 @@ func loadSingleExtension(path string) (*LoadedExtension, error) {
|
||||
return nil
|
||||
})
|
||||
},
|
||||
onStepStart: func(h func(StepStartEvent, Context)) {
|
||||
reg(StepStart, func(e Event, c Context) Result {
|
||||
h(e.(StepStartEvent), c)
|
||||
return nil
|
||||
})
|
||||
},
|
||||
onStepFinish: func(h func(StepFinishEvent, Context)) {
|
||||
reg(StepFinish, func(e Event, c Context) Result {
|
||||
h(e.(StepFinishEvent), c)
|
||||
return nil
|
||||
})
|
||||
},
|
||||
onReasoningStart: func(h func(ReasoningStartEvent, Context)) {
|
||||
reg(ReasoningStart, func(e Event, c Context) Result {
|
||||
h(e.(ReasoningStartEvent), c)
|
||||
return nil
|
||||
})
|
||||
},
|
||||
onWarnings: func(h func(WarningsEvent, Context)) {
|
||||
reg(Warnings, func(e Event, c Context) Result {
|
||||
h(e.(WarningsEvent), c)
|
||||
return nil
|
||||
})
|
||||
},
|
||||
onSource: func(h func(SourceEvent, Context)) {
|
||||
reg(Source, func(e Event, c Context) Result {
|
||||
h(e.(SourceEvent), c)
|
||||
return nil
|
||||
})
|
||||
},
|
||||
onError: func(h func(ErrorEvent, Context)) {
|
||||
reg(Error, func(e Event, c Context) Result {
|
||||
h(e.(ErrorEvent), c)
|
||||
return nil
|
||||
})
|
||||
},
|
||||
onRetry: func(h func(RetryEvent, Context)) {
|
||||
reg(Retry, func(e Event, c Context) Result {
|
||||
h(e.(RetryEvent), c)
|
||||
return nil
|
||||
})
|
||||
},
|
||||
onPrepareStep: func(h func(PrepareStepEvent, Context) *PrepareStepResult) {
|
||||
reg(PrepareStep, func(e Event, c Context) Result {
|
||||
r := h(e.(PrepareStepEvent), c)
|
||||
if r == nil {
|
||||
return nil
|
||||
}
|
||||
return *r
|
||||
})
|
||||
},
|
||||
}
|
||||
|
||||
// Call Init — the extension registers its handlers, tools, commands.
|
||||
|
||||
@@ -172,6 +172,17 @@ func Symbols() interp.Exports {
|
||||
"SessionStartEvent": reflect.ValueOf((*SessionStartEvent)(nil)),
|
||||
"SessionShutdownEvent": reflect.ValueOf((*SessionShutdownEvent)(nil)),
|
||||
"ModelChangeEvent": reflect.ValueOf((*ModelChangeEvent)(nil)),
|
||||
|
||||
// Step lifecycle events
|
||||
"StepStartEvent": reflect.ValueOf((*StepStartEvent)(nil)),
|
||||
"StepFinishEvent": reflect.ValueOf((*StepFinishEvent)(nil)),
|
||||
"ReasoningStartEvent": reflect.ValueOf((*ReasoningStartEvent)(nil)),
|
||||
"WarningsEvent": reflect.ValueOf((*WarningsEvent)(nil)),
|
||||
"SourceEvent": reflect.ValueOf((*SourceEvent)(nil)),
|
||||
"ErrorEvent": reflect.ValueOf((*ErrorEvent)(nil)),
|
||||
"RetryEvent": reflect.ValueOf((*RetryEvent)(nil)),
|
||||
"PrepareStepEvent": reflect.ValueOf((*PrepareStepEvent)(nil)),
|
||||
"PrepareStepResult": reflect.ValueOf((*PrepareStepResult)(nil)),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -63,6 +63,11 @@ type TreeManager struct {
|
||||
|
||||
// file is the open file handle for appending entries. Nil for in-memory.
|
||||
file *os.File
|
||||
|
||||
// writer is a buffered writer wrapping file. Writes go through this
|
||||
// buffer and are flushed to disk at explicit sync points (after each
|
||||
// public Append* call, in Close, etc.) to reduce syscall overhead.
|
||||
writer *bufio.Writer
|
||||
}
|
||||
|
||||
// --- Constructors ---
|
||||
@@ -105,11 +110,16 @@ func CreateTreeSession(cwd string) (*TreeManager, error) {
|
||||
return nil, fmt.Errorf("failed to create session file: %w", err)
|
||||
}
|
||||
tm.file = f
|
||||
tm.writer = bufio.NewWriter(f)
|
||||
|
||||
if err := tm.writeEntry(&header); err != nil {
|
||||
_ = f.Close()
|
||||
return nil, fmt.Errorf("failed to write session header: %w", err)
|
||||
}
|
||||
if err := tm.flushLocked(); err != nil {
|
||||
_ = f.Close()
|
||||
return nil, fmt.Errorf("failed to flush session header: %w", err)
|
||||
}
|
||||
|
||||
return tm, nil
|
||||
}
|
||||
@@ -150,6 +160,7 @@ func (tm *TreeManager) ForkToNewSession(cwd string, targetID string) (*TreeManag
|
||||
return nil, fmt.Errorf("failed to recreate session file: %w", err)
|
||||
}
|
||||
newTm.file = f
|
||||
newTm.writer = bufio.NewWriter(f)
|
||||
|
||||
if err := newTm.writeEntry(&newTm.header); err != nil {
|
||||
_ = f.Close()
|
||||
@@ -289,6 +300,12 @@ func (tm *TreeManager) ForkToNewSession(cwd string, targetID string) (*TreeManag
|
||||
}
|
||||
}
|
||||
|
||||
// Flush all buffered writes from the fork in a single syscall.
|
||||
if err := newTm.flushLocked(); err != nil {
|
||||
_ = f.Close()
|
||||
return nil, fmt.Errorf("failed to flush forked session: %w", err)
|
||||
}
|
||||
|
||||
// Set the leaf to the last entry in the new session.
|
||||
newTm.leafID = prevNewID
|
||||
|
||||
@@ -374,6 +391,7 @@ func OpenTreeSession(path string) (*TreeManager, error) {
|
||||
return nil, fmt.Errorf("failed to open session file for append: %w", err)
|
||||
}
|
||||
tm.file = f
|
||||
tm.writer = bufio.NewWriter(f)
|
||||
|
||||
return tm, nil
|
||||
}
|
||||
@@ -427,6 +445,9 @@ func (tm *TreeManager) AppendMessage(msg message.Message) (string, error) {
|
||||
if err := tm.appendAndPersist(entry); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := tm.flushLocked(); err != nil {
|
||||
return "", fmt.Errorf("failed to flush message: %w", err)
|
||||
}
|
||||
|
||||
tm.leafID = entry.ID
|
||||
return entry.ID, nil
|
||||
@@ -451,6 +472,9 @@ func (tm *TreeManager) AppendModelChange(provider, modelID string) (string, erro
|
||||
if err := tm.appendAndPersist(entry); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := tm.flushLocked(); err != nil {
|
||||
return "", fmt.Errorf("failed to flush model change: %w", err)
|
||||
}
|
||||
|
||||
tm.leafID = entry.ID
|
||||
return entry.ID, nil
|
||||
@@ -465,6 +489,9 @@ func (tm *TreeManager) AppendBranchSummary(fromID, summary string) (string, erro
|
||||
if err := tm.appendAndPersist(entry); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := tm.flushLocked(); err != nil {
|
||||
return "", fmt.Errorf("failed to flush branch summary: %w", err)
|
||||
}
|
||||
|
||||
tm.leafID = entry.ID
|
||||
return entry.ID, nil
|
||||
@@ -479,6 +506,9 @@ func (tm *TreeManager) AppendLabel(targetID, label string) (string, error) {
|
||||
if err := tm.appendAndPersist(entry); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := tm.flushLocked(); err != nil {
|
||||
return "", fmt.Errorf("failed to flush label: %w", err)
|
||||
}
|
||||
|
||||
tm.labels[targetID] = label
|
||||
tm.leafID = entry.ID
|
||||
@@ -494,6 +524,9 @@ func (tm *TreeManager) AppendSessionInfo(name string) (string, error) {
|
||||
if err := tm.appendAndPersist(entry); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := tm.flushLocked(); err != nil {
|
||||
return "", fmt.Errorf("failed to flush session info: %w", err)
|
||||
}
|
||||
|
||||
tm.sessionName = name
|
||||
tm.leafID = entry.ID
|
||||
@@ -510,6 +543,9 @@ func (tm *TreeManager) AppendExtensionData(extType, data string) (string, error)
|
||||
if err := tm.appendAndPersist(entry); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := tm.flushLocked(); err != nil {
|
||||
return "", fmt.Errorf("failed to flush extension data: %w", err)
|
||||
}
|
||||
|
||||
tm.leafID = entry.ID
|
||||
return entry.ID, nil
|
||||
@@ -541,6 +577,9 @@ func (tm *TreeManager) AppendCompaction(summary, firstKeptEntryID string, tokens
|
||||
if err := tm.appendAndPersist(entry); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := tm.flushLocked(); err != nil {
|
||||
return "", fmt.Errorf("failed to flush compaction: %w", err)
|
||||
}
|
||||
|
||||
tm.leafID = entry.ID
|
||||
return entry.ID, nil
|
||||
@@ -926,11 +965,31 @@ func (tm *TreeManager) IsEmpty() bool {
|
||||
return tm.MessageCount() == 0
|
||||
}
|
||||
|
||||
// Close closes the underlying file handle.
|
||||
// Flush writes any buffered data to the underlying file.
|
||||
func (tm *TreeManager) Flush() error {
|
||||
tm.mu.Lock()
|
||||
defer tm.mu.Unlock()
|
||||
return tm.flushLocked()
|
||||
}
|
||||
|
||||
// flushLocked writes buffered data to disk. Caller must hold the lock.
|
||||
func (tm *TreeManager) flushLocked() error {
|
||||
if tm.writer != nil {
|
||||
return tm.writer.Flush()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Close flushes any buffered writes and closes the underlying file handle.
|
||||
func (tm *TreeManager) Close() error {
|
||||
tm.mu.Lock()
|
||||
defer tm.mu.Unlock()
|
||||
if tm.file != nil {
|
||||
// Flush buffered data before closing.
|
||||
if tm.writer != nil {
|
||||
_ = tm.writer.Flush()
|
||||
tm.writer = nil
|
||||
}
|
||||
err := tm.file.Close()
|
||||
tm.file = nil
|
||||
return err
|
||||
@@ -1090,13 +1149,22 @@ func (tm *TreeManager) GetLastCompaction() *CompactionEntry {
|
||||
|
||||
// AddLLMMessages appends multiple LLM messages as entries. This is
|
||||
// used when syncing from the agent's ConversationMessages after a step.
|
||||
// All entries are buffered and flushed to disk in a single batch.
|
||||
func (tm *TreeManager) AddLLMMessages(msgs []fantasy.Message) error {
|
||||
tm.mu.Lock()
|
||||
defer tm.mu.Unlock()
|
||||
|
||||
for _, msg := range msgs {
|
||||
if _, err := tm.AppendLLMMessage(msg); err != nil {
|
||||
entry, err := NewMessageEntry(tm.leafID, message.FromLLMMessage(msg))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tm.appendAndPersist(entry); err != nil {
|
||||
return err
|
||||
}
|
||||
tm.leafID = entry.ID
|
||||
}
|
||||
return nil
|
||||
return tm.flushLocked()
|
||||
}
|
||||
|
||||
// Deprecated: Use AddLLMMessages instead.
|
||||
@@ -1148,12 +1216,20 @@ func (tm *TreeManager) appendAndPersist(entry any) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// writeEntry serializes an entry and appends it as a line to the file.
|
||||
// writeEntry serializes an entry and appends it to the buffered writer.
|
||||
// The data is not flushed to disk until flushLocked is called.
|
||||
func (tm *TreeManager) writeEntry(entry any) error {
|
||||
data, err := json.Marshal(entry)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal entry: %w", err)
|
||||
}
|
||||
if tm.writer != nil {
|
||||
if _, err := tm.writer.Write(data); err != nil {
|
||||
return err
|
||||
}
|
||||
return tm.writer.WriteByte('\n')
|
||||
}
|
||||
// Fallback for direct file writes (shouldn't happen in normal flow).
|
||||
data = append(data, '\n')
|
||||
_, err = tm.file.Write(data)
|
||||
return err
|
||||
|
||||
@@ -6,6 +6,8 @@ import (
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// FileSuggestion represents a single file, directory, or MCP resource
|
||||
@@ -31,6 +33,51 @@ type FileSuggestion struct {
|
||||
// maxFileSuggestions is the maximum number of file suggestions returned.
|
||||
const maxFileSuggestions = 20
|
||||
|
||||
// fileListCache caches the result of listFiles() keyed by directory to avoid
|
||||
// re-running git subprocesses on every keystroke during @file completion.
|
||||
var fileListCache struct {
|
||||
mu sync.Mutex
|
||||
dir string // searchDir that produced the cached entries
|
||||
cwd string // cwd used for the git query
|
||||
entries []FileSuggestion // cached file list
|
||||
expireAt time.Time // when the cache entry expires
|
||||
}
|
||||
|
||||
// fileListCacheTTL controls how long a cached file list stays valid.
|
||||
// During rapid typing the list is reused; after the TTL a fresh git
|
||||
// ls-files is executed so newly created files become visible.
|
||||
const fileListCacheTTL = 3 * time.Second
|
||||
|
||||
// getCachedFileList returns the file list for searchDir, using a short-lived
|
||||
// cache to avoid repeated subprocess calls during @file autocompletion.
|
||||
func getCachedFileList(searchDir, cwd string) []FileSuggestion {
|
||||
fileListCache.mu.Lock()
|
||||
defer fileListCache.mu.Unlock()
|
||||
|
||||
now := time.Now()
|
||||
if fileListCache.dir == searchDir &&
|
||||
fileListCache.cwd == cwd &&
|
||||
now.Before(fileListCache.expireAt) {
|
||||
// Return a copy so callers can mutate (e.g. prepend baseDir).
|
||||
cp := make([]FileSuggestion, len(fileListCache.entries))
|
||||
copy(cp, fileListCache.entries)
|
||||
return cp
|
||||
}
|
||||
|
||||
// Cache miss or expired — run the real (potentially expensive) lookup.
|
||||
files := listFiles(searchDir, cwd)
|
||||
|
||||
fileListCache.dir = searchDir
|
||||
fileListCache.cwd = cwd
|
||||
fileListCache.entries = files
|
||||
fileListCache.expireAt = now.Add(fileListCacheTTL)
|
||||
|
||||
// Return a copy.
|
||||
cp := make([]FileSuggestion, len(files))
|
||||
copy(cp, files)
|
||||
return cp
|
||||
}
|
||||
|
||||
// ExtractAtPrefix checks the current line for an @-file trigger at cursorCol.
|
||||
// It returns:
|
||||
// - hasAt: true if a valid @ trigger was found
|
||||
@@ -99,7 +146,7 @@ func GetFileSuggestions(prefix string, cwd string) []FileSuggestion {
|
||||
}
|
||||
}
|
||||
|
||||
files := listFiles(searchDir, cwd)
|
||||
files := getCachedFileList(searchDir, cwd)
|
||||
if len(files) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -109,8 +109,8 @@ func (m *TextMessageItem) renderContent(width int) string {
|
||||
// It accumulates content chunks and re-renders on each update for live display.
|
||||
type StreamingMessageItem struct {
|
||||
id string
|
||||
role string // "assistant" or "reasoning"
|
||||
content string // Accumulated streaming content
|
||||
role string // "assistant" or "reasoning"
|
||||
content strings.Builder // Accumulated streaming content
|
||||
timestamp time.Time
|
||||
startTime time.Time // When streaming started (for live duration counter)
|
||||
modelName string
|
||||
@@ -156,10 +156,10 @@ func (s *StreamingMessageItem) Render(width int) string {
|
||||
durationMs = time.Since(s.startTime).Milliseconds()
|
||||
}
|
||||
ty := createTypography(style.GetTheme())
|
||||
rendered = render.ReasoningBlock(s.content, durationMs, width, ty, style.GetTheme())
|
||||
rendered = render.ReasoningBlock(s.content.String(), durationMs, width, ty, style.GetTheme())
|
||||
} else {
|
||||
// Render as assistant message
|
||||
rendered = render.AssistantBlock(s.content, width, style.GetTheme())
|
||||
rendered = render.AssistantBlock(s.content.String(), width, style.GetTheme())
|
||||
}
|
||||
|
||||
// Cache and return (but reasoning is never cached due to live duration)
|
||||
@@ -187,7 +187,7 @@ func (s *StreamingMessageItem) Height() int {
|
||||
|
||||
// AppendChunk adds a content chunk and invalidates the render cache.
|
||||
func (s *StreamingMessageItem) AppendChunk(chunk string) {
|
||||
s.content += chunk
|
||||
s.content.WriteString(chunk)
|
||||
s.cachedWidth = 0 // Invalidate cache
|
||||
}
|
||||
|
||||
@@ -243,9 +243,7 @@ func (m *StreamingBashOutputItem) Render(width int) string {
|
||||
|
||||
// Header with command
|
||||
if m.command != "" {
|
||||
headerStyle := lipgloss.NewStyle().
|
||||
Foreground(theme.Muted).
|
||||
Italic(true)
|
||||
headerStyle := style.GetCachedStyles().BashHeader
|
||||
parts = append(parts, headerStyle.Render(fmt.Sprintf("▸ %s", m.command)))
|
||||
}
|
||||
|
||||
|
||||
@@ -338,7 +338,7 @@ func (r *MessageRenderer) RenderToolMessage(toolName, toolArgs, toolResult strin
|
||||
// Build the content: icon + name + params on first line, then body
|
||||
headerLine := styledIcon + " " + styledName
|
||||
if params != "" {
|
||||
headerLine += " " + lipgloss.NewStyle().Foreground(theme.Muted).Render(params)
|
||||
headerLine += " " + style.GetCachedStyles().ToolMuted.Render(params)
|
||||
}
|
||||
|
||||
// Get body content
|
||||
|
||||
@@ -45,7 +45,7 @@ func UserBlock(content string, width int, ty *herald.Typography, theme style.The
|
||||
// HighlightFileTokens wraps @file tokens in the given text with the theme
|
||||
// accent color so they stand out visually in rendered user messages.
|
||||
func HighlightFileTokens(text string, theme style.Theme) string {
|
||||
accentStyle := lipgloss.NewStyle().Foreground(theme.Accent).Bold(true)
|
||||
accentStyle := style.GetCachedStyles().FileTokenAccent
|
||||
return fileTokenPattern.ReplaceAllStringFunc(text, func(token string) string {
|
||||
return accentStyle.Render(token)
|
||||
})
|
||||
@@ -75,8 +75,8 @@ func ReasoningBlock(content string, duration int64, width int, ty *herald.Typogr
|
||||
if width > 4 {
|
||||
contentStr = wrapText(contentStr, width-4)
|
||||
}
|
||||
mutedStyle := lipgloss.NewStyle().Foreground(theme.Muted)
|
||||
contentRendered := mutedStyle.Render(ty.Italic(contentStr))
|
||||
cs := style.GetCachedStyles()
|
||||
contentRendered := cs.Muted.Render(ty.Italic(contentStr))
|
||||
|
||||
// Build label based on duration
|
||||
if duration > 0 {
|
||||
@@ -86,14 +86,14 @@ func ReasoningBlock(content string, duration int64, width int, ty *herald.Typogr
|
||||
} else {
|
||||
durationStr = fmt.Sprintf("%.1fs", float64(duration)/1000)
|
||||
}
|
||||
labelPart := lipgloss.NewStyle().Foreground(theme.VeryMuted).Render("Thought for ")
|
||||
durationPart := lipgloss.NewStyle().Foreground(theme.Accent).Render(durationStr)
|
||||
labelPart := cs.VeryMuted.Render("Thought for ")
|
||||
durationPart := cs.Accent.Render(durationStr)
|
||||
label := labelPart + durationPart
|
||||
rendered := contentRendered + "\n" + label
|
||||
return styleMarginBottom(theme, rendered)
|
||||
}
|
||||
|
||||
label := lipgloss.NewStyle().Foreground(theme.VeryMuted).Render("Thought")
|
||||
label := cs.VeryMuted.Render("Thought")
|
||||
rendered := contentRendered + "\n" + label
|
||||
|
||||
return styleMarginBottom(theme, rendered)
|
||||
@@ -194,7 +194,7 @@ func ToolBlock(displayName, params, body string, isError bool, width int, ty *he
|
||||
|
||||
// styleMarginBottom applies a 1-line margin bottom using the theme.
|
||||
func styleMarginBottom(theme style.Theme, content string) string {
|
||||
return lipgloss.NewStyle().MarginBottom(1).Render(content)
|
||||
return style.GetCachedStyles().MarginBottom1.Render(content)
|
||||
}
|
||||
|
||||
// wrapText soft-wraps a string to the given width using lipgloss, which is
|
||||
|
||||
+9
-11
@@ -21,12 +21,11 @@ func knightRiderFrames() []string {
|
||||
const numDots = 8
|
||||
const dot = "▪"
|
||||
|
||||
theme := style.GetTheme()
|
||||
|
||||
bright := lipgloss.NewStyle().Foreground(theme.Primary)
|
||||
med := lipgloss.NewStyle().Foreground(theme.Muted)
|
||||
dim := lipgloss.NewStyle().Foreground(theme.VeryMuted)
|
||||
off := lipgloss.NewStyle().Foreground(theme.MutedBorder)
|
||||
cs := style.GetCachedStyles()
|
||||
bright := cs.SpinnerBright
|
||||
med := cs.SpinnerMed
|
||||
dim := cs.SpinnerDim
|
||||
off := cs.SpinnerOff
|
||||
|
||||
// Scanner bounces: 0→7→0
|
||||
positions := make([]int, 0, 2*numDots-2)
|
||||
@@ -476,9 +475,8 @@ func (s *StreamComponent) renderReasoningBlock(reasoning string) string {
|
||||
if s.width > 4 {
|
||||
content = lipgloss.NewStyle().Width(s.width - 4).Render(content)
|
||||
}
|
||||
theme := GetTheme()
|
||||
mutedStyle := lipgloss.NewStyle().Foreground(theme.Muted)
|
||||
parts = append(parts, mutedStyle.Render(s.ty.Italic(content)))
|
||||
cs := style.GetCachedStyles()
|
||||
parts = append(parts, cs.Muted.Render(s.ty.Italic(content)))
|
||||
|
||||
// Duration footer with VeryMuted label and Accent duration.
|
||||
var duration time.Duration
|
||||
@@ -494,8 +492,8 @@ func (s *StreamComponent) renderReasoningBlock(reasoning string) string {
|
||||
} else {
|
||||
durationStr = fmt.Sprintf("%.1fs", duration.Seconds())
|
||||
}
|
||||
label := lipgloss.NewStyle().Foreground(theme.VeryMuted).Render("Thought for ")
|
||||
durationStyled := lipgloss.NewStyle().Foreground(theme.Accent).Render(durationStr)
|
||||
label := cs.VeryMuted.Render("Thought for ")
|
||||
durationStyled := cs.Accent.Render(durationStr)
|
||||
parts = append(parts, label+durationStyled)
|
||||
}
|
||||
|
||||
|
||||
@@ -40,6 +40,70 @@ func GetTheme() Theme {
|
||||
func SetTheme(theme Theme) {
|
||||
currentTheme = theme
|
||||
markdownTypographyCache = nil // invalidate cached renderer; colors may have changed
|
||||
styleCache = nil // invalidate cached styles; colors may have changed
|
||||
}
|
||||
|
||||
// CachedStyles holds pre-built lipgloss styles that are reused across
|
||||
// render frames. Invalidated by SetTheme, lazily rebuilt on next access.
|
||||
// Only accessed from BubbleTea's single-threaded Update/View cycle.
|
||||
type CachedStyles struct {
|
||||
// render/blocks.go
|
||||
FileTokenAccent lipgloss.Style // Foreground(Accent).Bold(true)
|
||||
Muted lipgloss.Style // Foreground(Muted)
|
||||
VeryMuted lipgloss.Style // Foreground(VeryMuted)
|
||||
Accent lipgloss.Style // Foreground(Accent)
|
||||
MarginBottom1 lipgloss.Style // MarginBottom(1)
|
||||
|
||||
// stream.go - spinner phases
|
||||
SpinnerBright lipgloss.Style // Foreground(Primary)
|
||||
SpinnerMed lipgloss.Style // Foreground(Muted)
|
||||
SpinnerDim lipgloss.Style // Foreground(VeryMuted)
|
||||
SpinnerOff lipgloss.Style // Foreground(MutedBorder)
|
||||
|
||||
// message_items.go - bash output
|
||||
BashHeader lipgloss.Style // Foreground(Muted).Italic(true)
|
||||
BashStderr lipgloss.Style // Foreground(Error)
|
||||
|
||||
// render/blocks.go - tool block
|
||||
ToolSuccess lipgloss.Style // Foreground(Success)
|
||||
ToolError lipgloss.Style // Foreground(Error)
|
||||
ToolInfo lipgloss.Style // Foreground(Info).Bold(true)
|
||||
ToolMuted lipgloss.Style // Foreground(Muted)
|
||||
|
||||
// common
|
||||
ErrorFg lipgloss.Style // Foreground(Error)
|
||||
TextBold lipgloss.Style // Foreground(Text).Bold(true)
|
||||
}
|
||||
|
||||
var styleCache *CachedStyles
|
||||
|
||||
// GetCachedStyles returns the pre-built style cache, creating it lazily
|
||||
// from the current theme. Invalidated by SetTheme.
|
||||
func GetCachedStyles() *CachedStyles {
|
||||
if styleCache != nil {
|
||||
return styleCache
|
||||
}
|
||||
theme := GetTheme()
|
||||
styleCache = &CachedStyles{
|
||||
FileTokenAccent: lipgloss.NewStyle().Foreground(theme.Accent).Bold(true),
|
||||
Muted: lipgloss.NewStyle().Foreground(theme.Muted),
|
||||
VeryMuted: lipgloss.NewStyle().Foreground(theme.VeryMuted),
|
||||
Accent: lipgloss.NewStyle().Foreground(theme.Accent),
|
||||
MarginBottom1: lipgloss.NewStyle().MarginBottom(1),
|
||||
SpinnerBright: lipgloss.NewStyle().Foreground(theme.Primary),
|
||||
SpinnerMed: lipgloss.NewStyle().Foreground(theme.Muted),
|
||||
SpinnerDim: lipgloss.NewStyle().Foreground(theme.VeryMuted),
|
||||
SpinnerOff: lipgloss.NewStyle().Foreground(theme.MutedBorder),
|
||||
BashHeader: lipgloss.NewStyle().Foreground(theme.Muted).Italic(true),
|
||||
BashStderr: lipgloss.NewStyle().Foreground(theme.Error),
|
||||
ToolSuccess: lipgloss.NewStyle().Foreground(theme.Success),
|
||||
ToolError: lipgloss.NewStyle().Foreground(theme.Error),
|
||||
ToolInfo: lipgloss.NewStyle().Foreground(theme.Info).Bold(true),
|
||||
ToolMuted: lipgloss.NewStyle().Foreground(theme.Muted),
|
||||
ErrorFg: lipgloss.NewStyle().Foreground(theme.Error),
|
||||
TextBold: lipgloss.NewStyle().Foreground(theme.Text).Bold(true),
|
||||
}
|
||||
return styleCache
|
||||
}
|
||||
|
||||
// MarkdownThemeColors defines colors for markdown rendering and syntax highlighting.
|
||||
|
||||
+1
-1
@@ -106,7 +106,7 @@ unsub2 := host.OnToolResult(func(e kit.ToolResultEvent) {
|
||||
})
|
||||
defer unsub2()
|
||||
|
||||
unsub3 := host.OnStreaming(func(e kit.MessageUpdateEvent) {
|
||||
unsub3 := host.OnMessageUpdate(func(e kit.MessageUpdateEvent) {
|
||||
fmt.Print(e.Chunk)
|
||||
})
|
||||
defer unsub3()
|
||||
|
||||
@@ -58,6 +58,31 @@ const (
|
||||
// EventSteerConsumed fires when one or more steering messages have been
|
||||
// injected into the agent turn via PrepareStep.
|
||||
EventSteerConsumed EventType = "steer_consumed"
|
||||
// EventStepStart fires when a new LLM call begins within a turn.
|
||||
EventStepStart EventType = "step_start"
|
||||
// EventStepFinish fires when a step completes, providing full step context
|
||||
// including whether tool calls were made, the finish reason, and usage stats.
|
||||
EventStepFinish EventType = "step_finish"
|
||||
// EventTextStart fires when the LLM begins generating text content.
|
||||
EventTextStart EventType = "text_start"
|
||||
// EventTextEnd fires when the LLM finishes generating text content.
|
||||
EventTextEnd EventType = "text_end"
|
||||
// EventReasoningStart fires when the LLM begins reasoning/thinking.
|
||||
EventReasoningStart EventType = "reasoning_start"
|
||||
// EventWarnings fires when the LLM provider returns warnings.
|
||||
EventWarnings EventType = "warnings"
|
||||
// EventSource fires when the LLM references a source (e.g. from web search).
|
||||
EventSource EventType = "source"
|
||||
// EventStreamFinish fires when a per-step LLM stream completes with
|
||||
// usage stats and a finish reason.
|
||||
EventStreamFinish EventType = "stream_finish"
|
||||
// EventError fires when an agent-level error occurs during streaming.
|
||||
// This is distinct from TurnEndEvent.Error — it fires at the point of
|
||||
// failure, before the turn ends.
|
||||
EventError EventType = "error"
|
||||
// EventRetry fires when the LLM provider request is retried after a
|
||||
// transient error.
|
||||
EventRetry EventType = "retry"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -379,6 +404,100 @@ type SteerConsumedEvent struct {
|
||||
// EventType implements Event.
|
||||
func (e SteerConsumedEvent) EventType() EventType { return EventSteerConsumed }
|
||||
|
||||
// StepStartEvent fires when a new LLM call begins within a multi-step agent turn.
|
||||
type StepStartEvent struct {
|
||||
StepNumber int
|
||||
}
|
||||
|
||||
// EventType implements Event.
|
||||
func (e StepStartEvent) EventType() EventType { return EventStepStart }
|
||||
|
||||
// StepFinishEvent fires when a step completes, providing full step context.
|
||||
// This is a unified event that carries the same data as the existing
|
||||
// ToolCallContentEvent and StepUsageEvent, plus additional step metadata.
|
||||
type StepFinishEvent struct {
|
||||
StepNumber int
|
||||
HasToolCalls bool
|
||||
FinishReason string
|
||||
Usage LLMUsage
|
||||
}
|
||||
|
||||
// EventType implements Event.
|
||||
func (e StepFinishEvent) EventType() EventType { return EventStepFinish }
|
||||
|
||||
// TextStartEvent fires when the LLM begins generating text content.
|
||||
// Paired with MessageUpdateEvent (deltas) and TextEndEvent.
|
||||
type TextStartEvent struct {
|
||||
ID string
|
||||
}
|
||||
|
||||
// EventType implements Event.
|
||||
func (e TextStartEvent) EventType() EventType { return EventTextStart }
|
||||
|
||||
// TextEndEvent fires when the LLM finishes generating text content.
|
||||
type TextEndEvent struct {
|
||||
ID string
|
||||
}
|
||||
|
||||
// EventType implements Event.
|
||||
func (e TextEndEvent) EventType() EventType { return EventTextEnd }
|
||||
|
||||
// ReasoningStartEvent fires when the LLM begins reasoning/thinking.
|
||||
// Paired with ReasoningDeltaEvent (deltas) and ReasoningCompleteEvent.
|
||||
type ReasoningStartEvent struct {
|
||||
ID string
|
||||
}
|
||||
|
||||
// EventType implements Event.
|
||||
func (e ReasoningStartEvent) EventType() EventType { return EventReasoningStart }
|
||||
|
||||
// WarningsEvent fires when the LLM provider returns warnings about the request.
|
||||
type WarningsEvent struct {
|
||||
Warnings []string
|
||||
}
|
||||
|
||||
// EventType implements Event.
|
||||
func (e WarningsEvent) EventType() EventType { return EventWarnings }
|
||||
|
||||
// SourceEvent fires when the LLM references a source (e.g. from web search tools).
|
||||
type SourceEvent struct {
|
||||
SourceType string
|
||||
ID string
|
||||
URL string
|
||||
Title string
|
||||
}
|
||||
|
||||
// EventType implements Event.
|
||||
func (e SourceEvent) EventType() EventType { return EventSource }
|
||||
|
||||
// StreamFinishEvent fires when a per-step LLM stream completes.
|
||||
// Provides per-stream usage stats and finish reason.
|
||||
type StreamFinishEvent struct {
|
||||
Usage LLMUsage
|
||||
FinishReason string
|
||||
}
|
||||
|
||||
// EventType implements Event.
|
||||
func (e StreamFinishEvent) EventType() EventType { return EventStreamFinish }
|
||||
|
||||
// ErrorEvent fires when an agent-level error occurs during streaming.
|
||||
// This is distinct from TurnEndEvent.Error — it fires at the point of failure.
|
||||
type ErrorEvent struct {
|
||||
Error error
|
||||
}
|
||||
|
||||
// EventType implements Event.
|
||||
func (e ErrorEvent) EventType() EventType { return EventError }
|
||||
|
||||
// RetryEvent fires when the LLM provider request is retried after a transient error.
|
||||
type RetryEvent struct {
|
||||
Attempt int
|
||||
Error error
|
||||
}
|
||||
|
||||
// EventType implements Event.
|
||||
func (e RetryEvent) EventType() EventType { return EventRetry }
|
||||
|
||||
// PasswordPromptEvent fires when a sudo command needs a password.
|
||||
// The TUI should display a password prompt and send the result back via ResponseCh.
|
||||
type PasswordPromptEvent struct {
|
||||
@@ -517,7 +636,16 @@ func (m *Kit) OnToolOutput(handler func(ToolOutputEvent)) func() {
|
||||
|
||||
// OnStreaming registers a handler that fires only for MessageUpdateEvent
|
||||
// (streaming text chunks). Returns an unsubscribe function.
|
||||
//
|
||||
// Deprecated: Use OnMessageUpdate instead. OnStreaming will be removed in a
|
||||
// future release.
|
||||
func (m *Kit) OnStreaming(handler func(MessageUpdateEvent)) func() {
|
||||
return m.OnMessageUpdate(handler)
|
||||
}
|
||||
|
||||
// OnMessageUpdate registers a handler that fires only for MessageUpdateEvent
|
||||
// (streaming text chunks). Returns an unsubscribe function.
|
||||
func (m *Kit) OnMessageUpdate(handler func(MessageUpdateEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if mu, ok := e.(MessageUpdateEvent); ok {
|
||||
handler(mu)
|
||||
@@ -555,6 +683,214 @@ func (m *Kit) OnTurnEnd(handler func(TurnEndEvent)) func() {
|
||||
})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Typed subscribers for previously unsubscribed event types
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// OnMessageStart registers a handler that fires only for MessageStartEvent.
|
||||
// Returns an unsubscribe function.
|
||||
func (m *Kit) OnMessageStart(handler func(MessageStartEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if ms, ok := e.(MessageStartEvent); ok {
|
||||
handler(ms)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// OnMessageEnd registers a handler that fires only for MessageEndEvent.
|
||||
// Returns an unsubscribe function.
|
||||
func (m *Kit) OnMessageEnd(handler func(MessageEndEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if me, ok := e.(MessageEndEvent); ok {
|
||||
handler(me)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// OnReasoningDelta registers a handler that fires only for ReasoningDeltaEvent.
|
||||
// Returns an unsubscribe function.
|
||||
func (m *Kit) OnReasoningDelta(handler func(ReasoningDeltaEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if rd, ok := e.(ReasoningDeltaEvent); ok {
|
||||
handler(rd)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// OnReasoningComplete registers a handler that fires only for ReasoningCompleteEvent.
|
||||
// Returns an unsubscribe function.
|
||||
func (m *Kit) OnReasoningComplete(handler func(ReasoningCompleteEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if rc, ok := e.(ReasoningCompleteEvent); ok {
|
||||
handler(rc)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// OnToolExecutionStart registers a handler that fires only for ToolExecutionStartEvent.
|
||||
// Returns an unsubscribe function.
|
||||
func (m *Kit) OnToolExecutionStart(handler func(ToolExecutionStartEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if tes, ok := e.(ToolExecutionStartEvent); ok {
|
||||
handler(tes)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// OnToolExecutionEnd registers a handler that fires only for ToolExecutionEndEvent.
|
||||
// Returns an unsubscribe function.
|
||||
func (m *Kit) OnToolExecutionEnd(handler func(ToolExecutionEndEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if tee, ok := e.(ToolExecutionEndEvent); ok {
|
||||
handler(tee)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// OnToolCallContent registers a handler that fires only for ToolCallContentEvent.
|
||||
// Returns an unsubscribe function.
|
||||
func (m *Kit) OnToolCallContent(handler func(ToolCallContentEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if tcc, ok := e.(ToolCallContentEvent); ok {
|
||||
handler(tcc)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// OnStepUsage registers a handler that fires only for StepUsageEvent.
|
||||
// Returns an unsubscribe function.
|
||||
func (m *Kit) OnStepUsage(handler func(StepUsageEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if su, ok := e.(StepUsageEvent); ok {
|
||||
handler(su)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// OnCompaction registers a handler that fires only for CompactionEvent.
|
||||
// Returns an unsubscribe function.
|
||||
func (m *Kit) OnCompaction(handler func(CompactionEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if ce, ok := e.(CompactionEvent); ok {
|
||||
handler(ce)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// OnSteerConsumed registers a handler that fires only for SteerConsumedEvent.
|
||||
// Returns an unsubscribe function.
|
||||
func (m *Kit) OnSteerConsumed(handler func(SteerConsumedEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if sc, ok := e.(SteerConsumedEvent); ok {
|
||||
handler(sc)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Typed subscribers for new event types
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// OnStepStart registers a handler that fires only for StepStartEvent.
|
||||
// Returns an unsubscribe function.
|
||||
func (m *Kit) OnStepStart(handler func(StepStartEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if ss, ok := e.(StepStartEvent); ok {
|
||||
handler(ss)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// OnStepFinish registers a handler that fires only for StepFinishEvent.
|
||||
// Returns an unsubscribe function.
|
||||
func (m *Kit) OnStepFinish(handler func(StepFinishEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if sf, ok := e.(StepFinishEvent); ok {
|
||||
handler(sf)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// OnTextStart registers a handler that fires only for TextStartEvent.
|
||||
// Returns an unsubscribe function.
|
||||
func (m *Kit) OnTextStart(handler func(TextStartEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if ts, ok := e.(TextStartEvent); ok {
|
||||
handler(ts)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// OnTextEnd registers a handler that fires only for TextEndEvent.
|
||||
// Returns an unsubscribe function.
|
||||
func (m *Kit) OnTextEnd(handler func(TextEndEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if te, ok := e.(TextEndEvent); ok {
|
||||
handler(te)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// OnReasoningStart registers a handler that fires only for ReasoningStartEvent.
|
||||
// Returns an unsubscribe function.
|
||||
func (m *Kit) OnReasoningStart(handler func(ReasoningStartEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if rs, ok := e.(ReasoningStartEvent); ok {
|
||||
handler(rs)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// OnWarnings registers a handler that fires only for WarningsEvent.
|
||||
// Returns an unsubscribe function.
|
||||
func (m *Kit) OnWarnings(handler func(WarningsEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if w, ok := e.(WarningsEvent); ok {
|
||||
handler(w)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// OnSource registers a handler that fires only for SourceEvent.
|
||||
// Returns an unsubscribe function.
|
||||
func (m *Kit) OnSource(handler func(SourceEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if s, ok := e.(SourceEvent); ok {
|
||||
handler(s)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// OnStreamFinish registers a handler that fires only for StreamFinishEvent.
|
||||
// Returns an unsubscribe function.
|
||||
func (m *Kit) OnStreamFinish(handler func(StreamFinishEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if sf, ok := e.(StreamFinishEvent); ok {
|
||||
handler(sf)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// OnError registers a handler that fires only for ErrorEvent.
|
||||
// Returns an unsubscribe function.
|
||||
func (m *Kit) OnError(handler func(ErrorEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if ee, ok := e.(ErrorEvent); ok {
|
||||
handler(ee)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// OnRetry registers a handler that fires only for RetryEvent.
|
||||
// Returns an unsubscribe function.
|
||||
func (m *Kit) OnRetry(handler func(RetryEvent)) func() {
|
||||
return m.Subscribe(func(e Event) {
|
||||
if r, ok := e.(RetryEvent); ok {
|
||||
handler(r)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Subagent event subscriptions
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package kit
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
@@ -190,6 +191,74 @@ func TestEventTypes(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewEventTypes verifies that each new event struct returns the correct EventType.
|
||||
func TestNewEventTypes(t *testing.T) {
|
||||
tests := []struct {
|
||||
event Event
|
||||
expected EventType
|
||||
}{
|
||||
{StepStartEvent{StepNumber: 0}, EventStepStart},
|
||||
{StepFinishEvent{StepNumber: 1, HasToolCalls: true}, EventStepFinish},
|
||||
{TextStartEvent{ID: "text-1"}, EventTextStart},
|
||||
{TextEndEvent{ID: "text-1"}, EventTextEnd},
|
||||
{ReasoningStartEvent{ID: "reason-1"}, EventReasoningStart},
|
||||
{WarningsEvent{Warnings: []string{"test"}}, EventWarnings},
|
||||
{SourceEvent{URL: "https://example.com", Title: "Example"}, EventSource},
|
||||
{StreamFinishEvent{FinishReason: "stop"}, EventStreamFinish},
|
||||
{ErrorEvent{Error: fmt.Errorf("test error")}, EventError},
|
||||
{RetryEvent{Attempt: 1, Error: fmt.Errorf("retry error")}, EventRetry},
|
||||
{ToolCallStartEvent{}, EventToolCallStart},
|
||||
{ToolCallDeltaEvent{}, EventToolCallDelta},
|
||||
{ToolCallEndEvent{}, EventToolCallEnd},
|
||||
{PasswordPromptEvent{}, EventPasswordPrompt},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
if got := tt.event.EventType(); got != tt.expected {
|
||||
t.Errorf("%T.EventType() = %q, want %q", tt.event, got, tt.expected)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewEventEmission verifies that new event types are properly emitted and received.
|
||||
func TestNewEventEmission(t *testing.T) {
|
||||
bus := newEventBus()
|
||||
var received []Event
|
||||
|
||||
bus.subscribe(func(e Event) {
|
||||
received = append(received, e)
|
||||
})
|
||||
|
||||
bus.emit(StepStartEvent{StepNumber: 0})
|
||||
bus.emit(TextStartEvent{ID: "text-1"})
|
||||
bus.emit(TextEndEvent{ID: "text-1"})
|
||||
bus.emit(ReasoningStartEvent{ID: "reason-1"})
|
||||
bus.emit(WarningsEvent{Warnings: []string{"low confidence"}})
|
||||
bus.emit(SourceEvent{URL: "https://example.com", Title: "Example"})
|
||||
bus.emit(StreamFinishEvent{FinishReason: "stop"})
|
||||
bus.emit(StepFinishEvent{StepNumber: 0, HasToolCalls: false, FinishReason: "stop"})
|
||||
bus.emit(ErrorEvent{Error: fmt.Errorf("test error")})
|
||||
bus.emit(RetryEvent{Attempt: 1, Error: fmt.Errorf("retry")})
|
||||
|
||||
if len(received) != 10 {
|
||||
t.Fatalf("expected 10 events, got %d", len(received))
|
||||
}
|
||||
|
||||
// Verify specific event fields
|
||||
if ss, ok := received[0].(StepStartEvent); !ok || ss.StepNumber != 0 {
|
||||
t.Errorf("event 0: expected StepStartEvent{StepNumber:0}, got %T %+v", received[0], received[0])
|
||||
}
|
||||
if ts, ok := received[1].(TextStartEvent); !ok || ts.ID != "text-1" {
|
||||
t.Errorf("event 1: expected TextStartEvent{ID:text-1}, got %T %+v", received[1], received[1])
|
||||
}
|
||||
if w, ok := received[4].(WarningsEvent); !ok || len(w.Warnings) != 1 || w.Warnings[0] != "low confidence" {
|
||||
t.Errorf("event 4: expected WarningsEvent with 1 warning, got %T %+v", received[4], received[4])
|
||||
}
|
||||
if sf, ok := received[7].(StepFinishEvent); !ok || sf.StepNumber != 0 || sf.HasToolCalls {
|
||||
t.Errorf("event 7: expected StepFinishEvent{StepNumber:0, HasToolCalls:false}, got %T %+v", received[7], received[7])
|
||||
}
|
||||
}
|
||||
|
||||
// TestEventBusListenerCanUnsubscribeInCallback verifies that a listener can
|
||||
// safely call its own unsubscribe function from within the callback.
|
||||
func TestEventBusListenerCanUnsubscribeInCallback(t *testing.T) {
|
||||
|
||||
@@ -356,4 +356,134 @@ func (m *Kit) bridgeExtensions(runner *extensions.Runner) {
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
// --- Step lifecycle observation events ---
|
||||
|
||||
if runner.HasHandlers(extensions.StepStart) {
|
||||
m.Subscribe(func(e Event) {
|
||||
if ev, ok := e.(StepStartEvent); ok {
|
||||
_, _ = runner.Emit(extensions.StepStartEvent{StepNumber: ev.StepNumber})
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
if runner.HasHandlers(extensions.StepFinish) {
|
||||
m.Subscribe(func(e Event) {
|
||||
if ev, ok := e.(StepFinishEvent); ok {
|
||||
_, _ = runner.Emit(extensions.StepFinishEvent{
|
||||
StepNumber: ev.StepNumber,
|
||||
HasToolCalls: ev.HasToolCalls,
|
||||
FinishReason: ev.FinishReason,
|
||||
InputTokens: ev.Usage.InputTokens,
|
||||
OutputTokens: ev.Usage.OutputTokens,
|
||||
CacheReadTokens: ev.Usage.CacheReadTokens,
|
||||
CacheWriteTokens: ev.Usage.CacheCreationTokens,
|
||||
})
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
if runner.HasHandlers(extensions.ReasoningStart) {
|
||||
m.Subscribe(func(e Event) {
|
||||
if ev, ok := e.(ReasoningStartEvent); ok {
|
||||
_, _ = runner.Emit(extensions.ReasoningStartEvent{ID: ev.ID})
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
if runner.HasHandlers(extensions.Warnings) {
|
||||
m.Subscribe(func(e Event) {
|
||||
if ev, ok := e.(WarningsEvent); ok {
|
||||
_, _ = runner.Emit(extensions.WarningsEvent{Warnings: ev.Warnings})
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
if runner.HasHandlers(extensions.Source) {
|
||||
m.Subscribe(func(e Event) {
|
||||
if ev, ok := e.(SourceEvent); ok {
|
||||
_, _ = runner.Emit(extensions.SourceEvent{
|
||||
SourceType: ev.SourceType,
|
||||
ID: ev.ID,
|
||||
URL: ev.URL,
|
||||
Title: ev.Title,
|
||||
})
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
if runner.HasHandlers(extensions.Error) {
|
||||
m.Subscribe(func(e Event) {
|
||||
if ev, ok := e.(ErrorEvent); ok {
|
||||
_, _ = runner.Emit(extensions.ErrorEvent{Error: ev.Error.Error()})
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
if runner.HasHandlers(extensions.Retry) {
|
||||
m.Subscribe(func(e Event) {
|
||||
if ev, ok := e.(RetryEvent); ok {
|
||||
_, _ = runner.Emit(extensions.RetryEvent{
|
||||
Attempt: ev.Attempt,
|
||||
Error: ev.Error.Error(),
|
||||
})
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// --- PrepareStep hook ---
|
||||
// Extension PrepareStep → SDK PrepareStep hook.
|
||||
// Same pattern as ContextPrepare: convert LLMMessage ↔ ContextMessage.
|
||||
if runner.HasHandlers(extensions.PrepareStep) {
|
||||
m.OnPrepareStep(HookPriorityNormal, func(h PrepareStepHook) *PrepareStepResult {
|
||||
// Convert LLM message slice to extension ContextMessage slice.
|
||||
extMsgs := make([]extensions.ContextMessage, len(h.Messages))
|
||||
for i, msg := range h.Messages {
|
||||
var sb strings.Builder
|
||||
for _, part := range msg.Content {
|
||||
if tp, ok := part.(LLMTextPart); ok {
|
||||
sb.WriteString(tp.Text)
|
||||
}
|
||||
}
|
||||
extMsgs[i] = extensions.ContextMessage{
|
||||
Index: i,
|
||||
Role: string(msg.Role),
|
||||
Content: sb.String(),
|
||||
}
|
||||
}
|
||||
|
||||
result, _ := runner.Emit(extensions.PrepareStepEvent{
|
||||
StepNumber: h.StepNumber,
|
||||
Messages: extMsgs,
|
||||
})
|
||||
r, ok := result.(extensions.PrepareStepResult)
|
||||
if !ok || r.Messages == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Rebuild LLM message slice from extension result.
|
||||
rebuilt := make([]LLMMessage, 0, len(r.Messages))
|
||||
for _, cm := range r.Messages {
|
||||
if cm.Index >= 0 && cm.Index < len(h.Messages) {
|
||||
rebuilt = append(rebuilt, h.Messages[cm.Index])
|
||||
} else {
|
||||
role := LLMRoleUser
|
||||
switch cm.Role {
|
||||
case "assistant":
|
||||
role = LLMRoleAssistant
|
||||
case "system":
|
||||
role = LLMRoleSystem
|
||||
case "tool":
|
||||
role = LLMRoleTool
|
||||
}
|
||||
rebuilt = append(rebuilt, LLMMessage{
|
||||
Role: role,
|
||||
Content: []LLMMessagePart{LLMTextPart{Text: cm.Content}},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return &PrepareStepResult{Messages: rebuilt}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -121,6 +121,32 @@ type BeforeCompactResult struct {
|
||||
Summary string
|
||||
}
|
||||
|
||||
// PrepareStepHook is the input for hooks that fire between steps within a
|
||||
// multi-step agent turn, with full message replacement capability. This is
|
||||
// the most powerful interception point — it fires after the existing steering
|
||||
// logic (if any) and before the messages are sent to the LLM.
|
||||
//
|
||||
// Use cases:
|
||||
// - Transforming tool results (e.g. converting image tool results to FilePart
|
||||
// user messages for vision models that don't support media in tool results)
|
||||
// - Dynamic tool filtering per step
|
||||
// - Mid-turn context injection beyond simple steering
|
||||
// - Custom stop conditions that inspect message history
|
||||
type PrepareStepHook struct {
|
||||
// StepNumber is the zero-based step index within the current turn.
|
||||
StepNumber int
|
||||
// Messages is the current context window that will be sent to the LLM.
|
||||
// This includes any steering messages already injected in this step.
|
||||
Messages []LLMMessage
|
||||
}
|
||||
|
||||
// PrepareStepResult can replace the context window between steps.
|
||||
type PrepareStepResult struct {
|
||||
// Messages replaces the entire context window for this step. If nil,
|
||||
// the original messages (including any steering) are used unchanged.
|
||||
Messages []LLMMessage
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Generic hook registry with priority ordering
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -248,6 +274,19 @@ func (m *Kit) OnBeforeCompact(p HookPriority, h func(BeforeCompactHook) *BeforeC
|
||||
return m.beforeCompact.register(p, h)
|
||||
}
|
||||
|
||||
// OnPrepareStep registers a hook that fires between steps within a multi-step
|
||||
// agent turn, after steering messages are injected and before the messages are
|
||||
// sent to the LLM. Return a non-nil PrepareStepResult with Messages to replace
|
||||
// the entire context window for this step. Hooks execute in priority order;
|
||||
// the first non-nil result wins. Returns an unregister function.
|
||||
//
|
||||
// This is the most powerful interception point in the agent lifecycle. It
|
||||
// enables patterns like transforming tool results, dynamic tool filtering,
|
||||
// and mid-turn context injection.
|
||||
func (m *Kit) OnPrepareStep(p HookPriority, h func(PrepareStepHook) *PrepareStepResult) func() {
|
||||
return m.prepareStep.register(p, h)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tool wrapping via hooks
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@@ -538,3 +538,75 @@ func TestKit_HookMethodsExist(t *testing.T) {
|
||||
u3()
|
||||
u4()
|
||||
}
|
||||
|
||||
// TestPrepareStepHookRegistry verifies registration and execution of PrepareStep hooks.
|
||||
func TestPrepareStepHookRegistry(t *testing.T) {
|
||||
hr := newHookRegistry[PrepareStepHook, PrepareStepResult]()
|
||||
|
||||
// Register a hook that appends a message.
|
||||
hr.register(HookPriorityNormal, func(h PrepareStepHook) *PrepareStepResult {
|
||||
if h.StepNumber == 0 {
|
||||
// On step 0, prepend a system message.
|
||||
newMsgs := make([]LLMMessage, 0, len(h.Messages)+1)
|
||||
newMsgs = append(newMsgs, fantasy.NewSystemMessage("injected"))
|
||||
newMsgs = append(newMsgs, h.Messages...)
|
||||
return &PrepareStepResult{Messages: newMsgs}
|
||||
}
|
||||
return nil // No modification for other steps.
|
||||
})
|
||||
|
||||
// Test step 0 — should modify messages.
|
||||
input := PrepareStepHook{
|
||||
StepNumber: 0,
|
||||
Messages: []LLMMessage{fantasy.NewUserMessage("hello")},
|
||||
}
|
||||
result := hr.run(input)
|
||||
if result == nil {
|
||||
t.Fatal("expected non-nil result for step 0")
|
||||
}
|
||||
if len(result.Messages) != 2 {
|
||||
t.Fatalf("expected 2 messages, got %d", len(result.Messages))
|
||||
}
|
||||
if result.Messages[0].Role != fantasy.MessageRoleSystem {
|
||||
t.Errorf("expected system message first, got role %q", result.Messages[0].Role)
|
||||
}
|
||||
|
||||
// Test step 1 — should return nil (no modification).
|
||||
input.StepNumber = 1
|
||||
result = hr.run(input)
|
||||
if result != nil {
|
||||
t.Errorf("expected nil result for step 1, got %+v", result)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPrepareStepHookPriority verifies that PrepareStep hooks respect priority ordering.
|
||||
func TestPrepareStepHookPriority(t *testing.T) {
|
||||
hr := newHookRegistry[PrepareStepHook, PrepareStepResult]()
|
||||
|
||||
var order []string
|
||||
|
||||
// Low priority — should run second.
|
||||
hr.register(HookPriorityLow, func(_ PrepareStepHook) *PrepareStepResult {
|
||||
order = append(order, "low")
|
||||
return nil
|
||||
})
|
||||
|
||||
// High priority — should run first and win.
|
||||
hr.register(HookPriorityHigh, func(h PrepareStepHook) *PrepareStepResult {
|
||||
order = append(order, "high")
|
||||
return &PrepareStepResult{Messages: h.Messages}
|
||||
})
|
||||
|
||||
input := PrepareStepHook{
|
||||
StepNumber: 0,
|
||||
Messages: []LLMMessage{fantasy.NewUserMessage("test")},
|
||||
}
|
||||
result := hr.run(input)
|
||||
|
||||
if result == nil {
|
||||
t.Fatal("expected non-nil result")
|
||||
}
|
||||
if len(order) != 1 || order[0] != "high" {
|
||||
t.Errorf("expected [high] (first non-nil wins), got %v", order)
|
||||
}
|
||||
}
|
||||
|
||||
+91
-27
@@ -66,6 +66,7 @@ type Kit struct {
|
||||
afterTurn *hookRegistry[AfterTurnHook, AfterTurnResult]
|
||||
contextPrepare *hookRegistry[ContextPrepareHook, ContextPrepareResult]
|
||||
beforeCompact *hookRegistry[BeforeCompactHook, BeforeCompactResult]
|
||||
prepareStep *hookRegistry[PrepareStepHook, PrepareStepResult]
|
||||
|
||||
// lastInputTokens stores the API-reported input token count from the
|
||||
// most recent turn. Used by GetContextStats() to return accurate usage
|
||||
@@ -1368,6 +1369,7 @@ func New(ctx context.Context, opts *Options) (*Kit, error) {
|
||||
afterTurn := newHookRegistry[AfterTurnHook, AfterTurnResult]()
|
||||
contextPrepare := newHookRegistry[ContextPrepareHook, ContextPrepareResult]()
|
||||
beforeCompact := newHookRegistry[BeforeCompactHook, BeforeCompactResult]()
|
||||
prepareStep := newHookRegistry[PrepareStepHook, PrepareStepResult]()
|
||||
|
||||
// Build agent setup options, pulling CLI-specific fields when available.
|
||||
// Pass the pre-built ProviderConfig and scalar viper snapshots so
|
||||
@@ -1461,6 +1463,7 @@ func New(ctx context.Context, opts *Options) (*Kit, error) {
|
||||
afterTurn: afterTurn,
|
||||
contextPrepare: contextPrepare,
|
||||
beforeCompact: beforeCompact,
|
||||
prepareStep: prepareStep,
|
||||
}
|
||||
|
||||
// Bridge extension events to SDK hooks.
|
||||
@@ -1781,12 +1784,19 @@ func (m *Kit) Subagent(ctx context.Context, cfg SubagentConfig) (*SubagentResult
|
||||
|
||||
// Create child Kit instance. Pass the parent's loaded MCP config to
|
||||
// avoid re-reading viper (which races with concurrent subagent spawns).
|
||||
// Streaming must be explicitly enabled — Options.Streaming defaults to
|
||||
// false, and New() unconditionally writes viper.Set("stream", opts.Streaming).
|
||||
// Without this, the subagent would (a) pollute viper global state for
|
||||
// other concurrent callers and (b) potentially hit provider-level
|
||||
// differences (e.g. Anthropic non-streaming timeouts with extended
|
||||
// thinking).
|
||||
childOpts := &Options{
|
||||
Model: model,
|
||||
SystemPrompt: systemPrompt,
|
||||
Tools: tools,
|
||||
NoSession: cfg.NoSession,
|
||||
Quiet: true,
|
||||
Streaming: true,
|
||||
MCPConfig: m.mcpConfig,
|
||||
}
|
||||
child, err := New(ctx, childOpts)
|
||||
@@ -1894,21 +1904,21 @@ func (m *Kit) generate(ctx context.Context, messages []fantasy.Message) (*agent.
|
||||
return sr, err
|
||||
})
|
||||
|
||||
return m.agent.GenerateWithLoopAndStreaming(ctx, messages,
|
||||
func(toolCallID, toolName, toolArgs string) {
|
||||
return m.agent.GenerateWithCallbacks(ctx, messages, agent.GenerateCallbacks{
|
||||
OnToolCall: func(toolCallID, toolName, toolArgs string) {
|
||||
m.events.emit(ToolCallEvent{
|
||||
ToolCallID: toolCallID, ToolName: toolName, ToolKind: toolKindFor(toolName),
|
||||
ToolArgs: toolArgs, ParsedArgs: parseToolArgs(toolArgs),
|
||||
})
|
||||
},
|
||||
func(toolCallID, toolName, toolArgs string, isStarting bool) {
|
||||
OnToolExecution: func(toolCallID, toolName, toolArgs string, isStarting bool) {
|
||||
if isStarting {
|
||||
m.events.emit(ToolExecutionStartEvent{ToolCallID: toolCallID, ToolName: toolName, ToolKind: toolKindFor(toolName), ToolArgs: toolArgs})
|
||||
} else {
|
||||
m.events.emit(ToolExecutionEndEvent{ToolCallID: toolCallID, ToolName: toolName, ToolKind: toolKindFor(toolName)})
|
||||
}
|
||||
},
|
||||
func(toolCallID, toolName, toolArgs, resultText, metadata string, isError bool) {
|
||||
OnToolResult: func(toolCallID, toolName, toolArgs, resultText, metadata string, isError bool) {
|
||||
evt := ToolResultEvent{
|
||||
ToolCallID: toolCallID, ToolName: toolName, ToolKind: toolKindFor(toolName),
|
||||
ToolArgs: toolArgs, ParsedArgs: parseToolArgs(toolArgs),
|
||||
@@ -1922,17 +1932,17 @@ func (m *Kit) generate(ctx context.Context, messages []fantasy.Message) (*agent.
|
||||
}
|
||||
m.events.emit(evt)
|
||||
},
|
||||
func(content string) {
|
||||
OnResponse: func(content string) {
|
||||
m.events.emit(ResponseEvent{Content: content})
|
||||
},
|
||||
func(content string) {
|
||||
OnToolCallContent: func(content string) {
|
||||
m.events.emit(ToolCallContentEvent{Content: content})
|
||||
},
|
||||
// <think> tag filtering: models like Qwen/DeepSeek wrap reasoning inside
|
||||
// <think>...</think> tags in the regular text stream. We intercept those
|
||||
// spans here and re-route them as ReasoningDeltaEvent/ReasoningCompleteEvent
|
||||
// so callers always receive clean, tag-free text and structured reasoning.
|
||||
func() func(chunk string) {
|
||||
OnStreamingResponse: func() func(chunk string) {
|
||||
const (
|
||||
thinkOpen = "<think>"
|
||||
thinkClose = "</think>"
|
||||
@@ -1968,14 +1978,13 @@ func (m *Kit) generate(ctx context.Context, messages []fantasy.Message) (*agent.
|
||||
}
|
||||
}
|
||||
}(),
|
||||
func(delta string) {
|
||||
OnReasoningDelta: func(delta string) {
|
||||
m.events.emit(ReasoningDeltaEvent{Delta: delta})
|
||||
},
|
||||
func() {
|
||||
OnReasoningComplete: func() {
|
||||
m.events.emit(ReasoningCompleteEvent{})
|
||||
},
|
||||
func(toolCallID, toolName, chunk string, isStderr bool) {
|
||||
// Emit tool output chunk event for streaming bash output
|
||||
OnToolOutput: func(toolCallID, toolName, chunk string, isStderr bool) {
|
||||
m.events.emit(ToolOutputEvent{
|
||||
ToolCallID: toolCallID,
|
||||
ToolName: toolName,
|
||||
@@ -1984,18 +1993,13 @@ func (m *Kit) generate(ctx context.Context, messages []fantasy.Message) (*agent.
|
||||
})
|
||||
},
|
||||
// Persist step messages incrementally so that progress survives
|
||||
// crashes and long-running turns don't lose work. Each step's
|
||||
// messages are persisted as a unit: for tool-calling steps this is
|
||||
// the assistant message (with tool_use parts) + tool-role message
|
||||
// (with tool_result parts) as a pair; for the final step it's the
|
||||
// assistant text/reasoning message alone.
|
||||
func(stepMessages []fantasy.Message) {
|
||||
// crashes and long-running turns don't lose work.
|
||||
OnStepMessages: func(stepMessages []fantasy.Message) {
|
||||
for _, msg := range stepMessages {
|
||||
_, _ = m.session.AppendMessage(msg)
|
||||
}
|
||||
},
|
||||
func(inputTokens, outputTokens, cacheReadTokens, cacheCreationTokens int64) {
|
||||
// Emit step usage event for real-time cost tracking
|
||||
OnStepUsage: func(inputTokens, outputTokens, cacheReadTokens, cacheCreationTokens int64) {
|
||||
if viper.GetBool("debug") {
|
||||
log.Printf("DEBUG Kit.generate emitting StepUsageEvent: input=%d output=%d cacheRead=%d cacheCreate=%d",
|
||||
inputTokens, outputTokens, cacheReadTokens, cacheCreationTokens,
|
||||
@@ -2009,37 +2013,97 @@ func (m *Kit) generate(ctx context.Context, messages []fantasy.Message) (*agent.
|
||||
})
|
||||
},
|
||||
// Password prompt handler for sudo commands
|
||||
func(prompt string) (string, bool) {
|
||||
// Emit event to TUI and wait for response via channel
|
||||
OnPasswordPrompt: func(prompt string) (string, bool) {
|
||||
responseCh := make(chan PasswordPromptResponse, 1)
|
||||
m.events.emit(PasswordPromptEvent{
|
||||
Prompt: prompt,
|
||||
ResponseCh: responseCh,
|
||||
})
|
||||
// Wait for response (TUI will send password or cancel)
|
||||
resp := <-responseCh
|
||||
return resp.Password, resp.Cancelled
|
||||
},
|
||||
// Tool call argument streaming — fire as the LLM generates tool arguments
|
||||
func(toolCallID, toolName string) {
|
||||
// Tool call argument streaming
|
||||
OnToolCallStart: func(toolCallID, toolName string) {
|
||||
m.events.emit(ToolCallStartEvent{
|
||||
ToolCallID: toolCallID,
|
||||
ToolName: toolName,
|
||||
ToolKind: toolKindFor(toolName),
|
||||
})
|
||||
},
|
||||
func(toolCallID, delta string) {
|
||||
OnToolCallDelta: func(toolCallID, delta string) {
|
||||
m.events.emit(ToolCallDeltaEvent{
|
||||
ToolCallID: toolCallID,
|
||||
Delta: delta,
|
||||
})
|
||||
},
|
||||
func(toolCallID string) {
|
||||
OnToolCallEnd: func(toolCallID string) {
|
||||
m.events.emit(ToolCallEndEvent{
|
||||
ToolCallID: toolCallID,
|
||||
})
|
||||
},
|
||||
)
|
||||
|
||||
// New callbacks for previously unwired Fantasy lifecycle events.
|
||||
OnStepStart: func(stepNumber int) {
|
||||
m.events.emit(StepStartEvent{StepNumber: stepNumber})
|
||||
},
|
||||
OnStepFinish: func(stepNumber int, hasToolCalls bool, finishReason string, usage fantasy.Usage) {
|
||||
m.events.emit(StepFinishEvent{
|
||||
StepNumber: stepNumber,
|
||||
HasToolCalls: hasToolCalls,
|
||||
FinishReason: finishReason,
|
||||
Usage: usage,
|
||||
})
|
||||
},
|
||||
OnTextStart: func(id string) {
|
||||
m.events.emit(TextStartEvent{ID: id})
|
||||
},
|
||||
OnTextEnd: func(id string) {
|
||||
m.events.emit(TextEndEvent{ID: id})
|
||||
},
|
||||
OnReasoningStart: func(id string) {
|
||||
m.events.emit(ReasoningStartEvent{ID: id})
|
||||
},
|
||||
OnWarnings: func(warnings []string) {
|
||||
m.events.emit(WarningsEvent{Warnings: warnings})
|
||||
},
|
||||
OnSource: func(sourceType, id, url, title string) {
|
||||
m.events.emit(SourceEvent{
|
||||
SourceType: sourceType,
|
||||
ID: id,
|
||||
URL: url,
|
||||
Title: title,
|
||||
})
|
||||
},
|
||||
OnStreamFinish: func(usage fantasy.Usage, finishReason string) {
|
||||
m.events.emit(StreamFinishEvent{
|
||||
Usage: usage,
|
||||
FinishReason: finishReason,
|
||||
})
|
||||
},
|
||||
OnError: func(err error) {
|
||||
m.events.emit(ErrorEvent{Error: err})
|
||||
},
|
||||
OnRetry: func(attempt int, err error) {
|
||||
m.events.emit(RetryEvent{Attempt: attempt, Error: err})
|
||||
},
|
||||
// PrepareStep hook — compose with steering (handled in agent layer)
|
||||
// and then run SDK consumer hooks.
|
||||
OnPrepareStep: func() agent.PrepareStepHandler {
|
||||
if !m.prepareStep.hasHooks() {
|
||||
return nil
|
||||
}
|
||||
return func(stepNumber int, messages []fantasy.Message) []fantasy.Message {
|
||||
hookResult := m.prepareStep.run(PrepareStepHook{
|
||||
StepNumber: stepNumber,
|
||||
Messages: messages,
|
||||
})
|
||||
if hookResult != nil && hookResult.Messages != nil {
|
||||
return hookResult.Messages
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}(),
|
||||
})
|
||||
}
|
||||
|
||||
// runTurn is the shared lifecycle for every prompt mode:
|
||||
|
||||
+52
-22
@@ -2,6 +2,7 @@ package kit
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"charm.land/fantasy"
|
||||
|
||||
@@ -52,6 +53,22 @@ func ErrorResult(content string) ToolOutput {
|
||||
return ToolOutput{Content: content, IsError: true}
|
||||
}
|
||||
|
||||
// ImageResult creates a [ToolOutput] that returns an image to the LLM.
|
||||
// The data is the raw image bytes and mediaType is the MIME type
|
||||
// (e.g. "image/png", "image/jpeg"). The optional text content accompanies
|
||||
// the image and is visible to the LLM alongside it.
|
||||
func ImageResult(content string, data []byte, mediaType string) ToolOutput {
|
||||
return ToolOutput{Content: content, Data: data, MediaType: mediaType}
|
||||
}
|
||||
|
||||
// MediaResult creates a [ToolOutput] that returns non-image binary media
|
||||
// (e.g. audio, video) to the LLM. The data is the raw bytes and mediaType
|
||||
// is the MIME type (e.g. "audio/wav", "video/mp4"). The optional text
|
||||
// content accompanies the media.
|
||||
func MediaResult(content string, data []byte, mediaType string) ToolOutput {
|
||||
return ToolOutput{Content: content, Data: data, MediaType: mediaType}
|
||||
}
|
||||
|
||||
// toolCallIDKey is the context key for the tool call ID.
|
||||
type toolCallIDKey struct{}
|
||||
|
||||
@@ -63,9 +80,35 @@ func ToolCallIDFromContext(ctx context.Context) string {
|
||||
return s
|
||||
}
|
||||
|
||||
// toolOutputToResponse converts a [ToolOutput] into the underlying
|
||||
// framework's ToolResponse, inferring the response Type from Data/MediaType
|
||||
// so that binary content (images, audio, etc.) is forwarded to the LLM
|
||||
// instead of being silently dropped.
|
||||
func toolOutputToResponse(result ToolOutput) fantasy.ToolResponse {
|
||||
resp := fantasy.ToolResponse{
|
||||
Content: result.Content,
|
||||
IsError: result.IsError,
|
||||
Data: result.Data,
|
||||
MediaType: result.MediaType,
|
||||
}
|
||||
// Infer response type from binary data so the downstream framework
|
||||
// creates a media content block instead of a plain-text one.
|
||||
if len(result.Data) > 0 && result.MediaType != "" {
|
||||
if strings.HasPrefix(result.MediaType, "image/") {
|
||||
resp.Type = "image"
|
||||
} else {
|
||||
resp.Type = "media"
|
||||
}
|
||||
}
|
||||
if result.Metadata != nil {
|
||||
resp = fantasy.WithResponseMetadata(resp, result.Metadata)
|
||||
}
|
||||
return resp
|
||||
}
|
||||
|
||||
// NewTool creates a custom [Tool] with automatic JSON schema generation from
|
||||
// the TInput struct type. The handler receives a typed input (deserialized
|
||||
// from the LLM's JSON arguments) and returns a [ToolResult].
|
||||
// from the LLM's JSON arguments) and returns a [ToolOutput].
|
||||
//
|
||||
// Struct tags on TInput control the generated schema:
|
||||
//
|
||||
@@ -77,6 +120,11 @@ func ToolCallIDFromContext(ctx context.Context) string {
|
||||
// The tool call ID is injected into the context and can be retrieved with
|
||||
// [ToolCallIDFromContext].
|
||||
//
|
||||
// Binary results: When [ToolOutput.Data] and [ToolOutput.MediaType] are set,
|
||||
// the response type is automatically inferred so the LLM receives the binary
|
||||
// content (e.g. an image) instead of only the text. Use [ImageResult] or
|
||||
// [MediaResult] for convenience.
|
||||
//
|
||||
// Example:
|
||||
//
|
||||
// type WeatherInput struct {
|
||||
@@ -84,7 +132,7 @@ func ToolCallIDFromContext(ctx context.Context) string {
|
||||
// }
|
||||
//
|
||||
// tool := kit.NewTool("get_weather", "Get weather for a city",
|
||||
// func(ctx context.Context, input WeatherInput) (kit.ToolResult, error) {
|
||||
// func(ctx context.Context, input WeatherInput) (kit.ToolOutput, error) {
|
||||
// return kit.TextResult("72°F, sunny in " + input.City), nil
|
||||
// },
|
||||
// )
|
||||
@@ -96,16 +144,7 @@ func NewTool[TInput any](name, description string, fn func(ctx context.Context,
|
||||
if err != nil {
|
||||
return fantasy.NewTextErrorResponse(err.Error()), nil
|
||||
}
|
||||
resp := fantasy.ToolResponse{
|
||||
Content: result.Content,
|
||||
IsError: result.IsError,
|
||||
Data: result.Data,
|
||||
MediaType: result.MediaType,
|
||||
}
|
||||
if result.Metadata != nil {
|
||||
resp = fantasy.WithResponseMetadata(resp, result.Metadata)
|
||||
}
|
||||
return resp, nil
|
||||
return toolOutputToResponse(result), nil
|
||||
},
|
||||
)
|
||||
}
|
||||
@@ -121,16 +160,7 @@ func NewParallelTool[TInput any](name, description string, fn func(ctx context.C
|
||||
if err != nil {
|
||||
return fantasy.NewTextErrorResponse(err.Error()), nil
|
||||
}
|
||||
resp := fantasy.ToolResponse{
|
||||
Content: result.Content,
|
||||
IsError: result.IsError,
|
||||
Data: result.Data,
|
||||
MediaType: result.MediaType,
|
||||
}
|
||||
if result.Metadata != nil {
|
||||
resp = fantasy.WithResponseMetadata(resp, result.Metadata)
|
||||
}
|
||||
return resp, nil
|
||||
return toolOutputToResponse(result), nil
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
@@ -117,3 +117,149 @@ func TestToolOutput_BinaryData(t *testing.T) {
|
||||
t.Errorf("MediaType = %q, want %q", r.MediaType, "image/png")
|
||||
}
|
||||
}
|
||||
|
||||
// TestImageResult verifies the ImageResult convenience constructor.
|
||||
func TestImageResult(t *testing.T) {
|
||||
data := []byte{0x89, 0x50, 0x4E, 0x47}
|
||||
r := kit.ImageResult("here is the image", data, "image/png")
|
||||
if r.Content != "here is the image" {
|
||||
t.Errorf("Content = %q, want %q", r.Content, "here is the image")
|
||||
}
|
||||
if len(r.Data) != 4 {
|
||||
t.Errorf("Data len = %d, want 4", len(r.Data))
|
||||
}
|
||||
if r.MediaType != "image/png" {
|
||||
t.Errorf("MediaType = %q, want %q", r.MediaType, "image/png")
|
||||
}
|
||||
if r.IsError {
|
||||
t.Error("ImageResult should not set IsError")
|
||||
}
|
||||
}
|
||||
|
||||
// TestMediaResult verifies the MediaResult convenience constructor.
|
||||
func TestMediaResult(t *testing.T) {
|
||||
data := []byte{0xFF, 0xFB, 0x90, 0x00}
|
||||
r := kit.MediaResult("audio clip", data, "audio/mpeg")
|
||||
if r.Content != "audio clip" {
|
||||
t.Errorf("Content = %q, want %q", r.Content, "audio clip")
|
||||
}
|
||||
if len(r.Data) != 4 {
|
||||
t.Errorf("Data len = %d, want 4", len(r.Data))
|
||||
}
|
||||
if r.MediaType != "audio/mpeg" {
|
||||
t.Errorf("MediaType = %q, want %q", r.MediaType, "audio/mpeg")
|
||||
}
|
||||
if r.IsError {
|
||||
t.Error("MediaResult should not set IsError")
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewTool_BinaryImageResponse verifies that NewTool correctly infers the
|
||||
// response type for image data so binary content is forwarded to the LLM
|
||||
// (issue #17).
|
||||
func TestNewTool_BinaryImageResponse(t *testing.T) {
|
||||
type Input struct {
|
||||
Path string `json:"path"`
|
||||
}
|
||||
|
||||
imgData := []byte{0x89, 0x50, 0x4E, 0x47} // PNG magic bytes
|
||||
|
||||
tool := kit.NewTool("read_image", "Read an image file",
|
||||
func(ctx context.Context, input Input) (kit.ToolOutput, error) {
|
||||
return kit.ImageResult("Here is the image", imgData, "image/png"), nil
|
||||
},
|
||||
)
|
||||
|
||||
// Run the tool and inspect the raw ToolResponse via the AgentTool interface.
|
||||
resp, err := tool.Run(context.Background(), kit.LLMToolCall{
|
||||
ID: "call_1",
|
||||
Name: "read_image",
|
||||
Input: `{"path": "test.png"}`,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Run() error: %v", err)
|
||||
}
|
||||
|
||||
// The Type field must be "image" so the downstream framework creates a
|
||||
// media content block instead of discarding the binary data.
|
||||
if resp.Type != "image" {
|
||||
t.Errorf("ToolResponse.Type = %q, want %q", resp.Type, "image")
|
||||
}
|
||||
if len(resp.Data) != 4 {
|
||||
t.Errorf("ToolResponse.Data len = %d, want 4", len(resp.Data))
|
||||
}
|
||||
if resp.MediaType != "image/png" {
|
||||
t.Errorf("ToolResponse.MediaType = %q, want %q", resp.MediaType, "image/png")
|
||||
}
|
||||
if resp.Content != "Here is the image" {
|
||||
t.Errorf("ToolResponse.Content = %q, want %q", resp.Content, "Here is the image")
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewTool_BinaryMediaResponse verifies type inference for non-image media.
|
||||
func TestNewTool_BinaryMediaResponse(t *testing.T) {
|
||||
type Input struct{}
|
||||
|
||||
tool := kit.NewTool("get_audio", "Get audio",
|
||||
func(ctx context.Context, input Input) (kit.ToolOutput, error) {
|
||||
return kit.MediaResult("audio clip", []byte{0xFF, 0xFB}, "audio/mpeg"), nil
|
||||
},
|
||||
)
|
||||
|
||||
resp, err := tool.Run(context.Background(), kit.LLMToolCall{
|
||||
ID: "call_2",
|
||||
Name: "get_audio",
|
||||
Input: `{}`,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Run() error: %v", err)
|
||||
}
|
||||
if resp.Type != "media" {
|
||||
t.Errorf("ToolResponse.Type = %q, want %q", resp.Type, "media")
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewTool_TextResponseTypeNotSet verifies that text-only responses do NOT
|
||||
// get an inferred type (preserving existing behavior).
|
||||
func TestNewTool_TextResponseTypeNotSet(t *testing.T) {
|
||||
type Input struct{}
|
||||
|
||||
tool := kit.NewTool("echo", "Echo",
|
||||
func(ctx context.Context, input Input) (kit.ToolOutput, error) {
|
||||
return kit.TextResult("hello"), nil
|
||||
},
|
||||
)
|
||||
|
||||
resp, err := tool.Run(context.Background(), kit.LLMToolCall{
|
||||
ID: "call_3", Name: "echo", Input: `{}`,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Run() error: %v", err)
|
||||
}
|
||||
// Text responses should not have Type set (the framework treats "" as text).
|
||||
if resp.Type != "" {
|
||||
t.Errorf("ToolResponse.Type = %q, want empty string for text responses", resp.Type)
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewParallelTool_BinaryImageResponse mirrors the NewTool binary test for
|
||||
// NewParallelTool.
|
||||
func TestNewParallelTool_BinaryImageResponse(t *testing.T) {
|
||||
type Input struct{}
|
||||
|
||||
tool := kit.NewParallelTool("snap", "Take a snapshot",
|
||||
func(ctx context.Context, input Input) (kit.ToolOutput, error) {
|
||||
return kit.ImageResult("snapshot", []byte{0xFF, 0xD8}, "image/jpeg"), nil
|
||||
},
|
||||
)
|
||||
|
||||
resp, err := tool.Run(context.Background(), kit.LLMToolCall{
|
||||
ID: "call_4", Name: "snap", Input: `{}`,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Run() error: %v", err)
|
||||
}
|
||||
if resp.Type != "image" {
|
||||
t.Errorf("ToolResponse.Type = %q, want %q", resp.Type, "image")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -157,6 +157,18 @@ type LLMTextPart = fantasy.TextPart
|
||||
// LLMReasoningPart is a reasoning/chain-of-thought content part.
|
||||
type LLMReasoningPart = fantasy.ReasoningPart
|
||||
|
||||
// LLMToolCall represents the raw tool invocation passed to a [Tool]'s Run
|
||||
// method. It carries the call ID, tool name, and the JSON-encoded input
|
||||
// arguments from the LLM. This is the execution-layer call object — distinct
|
||||
// from [ToolCall] (a message content part).
|
||||
type LLMToolCall = fantasy.ToolCall
|
||||
|
||||
// LLMToolResponse represents the raw response returned from a [Tool]'s Run
|
||||
// method. Most SDK consumers should use [ToolOutput] with [NewTool] /
|
||||
// [NewParallelTool] instead — this alias is provided for advanced use cases
|
||||
// that need to call Tool.Run() directly (e.g. testing).
|
||||
type LLMToolResponse = fantasy.ToolResponse
|
||||
|
||||
// LLMToolCallPart represents an LLM-initiated tool invocation within a message.
|
||||
type LLMToolCallPart = fantasy.ToolCallPart
|
||||
|
||||
|
||||
+86
-2
@@ -281,7 +281,7 @@ host.OnToolOutput(func(e kit.ToolOutputEvent) {
|
||||
// Streaming bash output chunks
|
||||
})
|
||||
|
||||
host.OnStreaming(func(e kit.MessageUpdateEvent) {
|
||||
host.OnMessageUpdate(func(e kit.MessageUpdateEvent) {
|
||||
fmt.Print(e.Chunk) // real-time text streaming
|
||||
})
|
||||
|
||||
@@ -296,8 +296,64 @@ host.OnTurnStart(func(e kit.TurnStartEvent) {
|
||||
host.OnTurnEnd(func(e kit.TurnEndEvent) {
|
||||
// e.Response, e.Error, e.StopReason
|
||||
})
|
||||
|
||||
host.OnStepStart(func(e kit.StepStartEvent) {
|
||||
// e.StepNumber — which LLM call step (1-based)
|
||||
})
|
||||
|
||||
host.OnStepFinish(func(e kit.StepFinishEvent) {
|
||||
// e.StepNumber, e.HasToolCalls, e.FinishReason, e.Usage (LLMUsage)
|
||||
})
|
||||
|
||||
host.OnWarnings(func(e kit.WarningsEvent) {
|
||||
for _, w := range e.Warnings {
|
||||
log.Printf("warning: %s", w)
|
||||
}
|
||||
})
|
||||
|
||||
host.OnError(func(e kit.ErrorEvent) {
|
||||
log.Printf("agent error: %v", e.Error)
|
||||
})
|
||||
|
||||
host.OnRetry(func(e kit.RetryEvent) {
|
||||
log.Printf("retrying (attempt %d): %v", e.Attempt, e.Error)
|
||||
})
|
||||
|
||||
host.OnTextStart(func(e kit.TextStartEvent) {
|
||||
// e.ID — content block ID
|
||||
})
|
||||
|
||||
host.OnTextEnd(func(e kit.TextEndEvent) {
|
||||
// e.ID — content block ID
|
||||
})
|
||||
|
||||
host.OnReasoningStart(func(e kit.ReasoningStartEvent) {
|
||||
// e.ID — reasoning block ID
|
||||
})
|
||||
|
||||
host.OnSource(func(e kit.SourceEvent) {
|
||||
// e.SourceType, e.ID, e.URL, e.Title
|
||||
})
|
||||
|
||||
host.OnStreamFinish(func(e kit.StreamFinishEvent) {
|
||||
// e.Usage (LLMUsage), e.FinishReason
|
||||
})
|
||||
|
||||
// Additional typed subscribers for previously generic-only events:
|
||||
host.OnMessageStart(func(e kit.MessageStartEvent) {})
|
||||
host.OnMessageEnd(func(e kit.MessageEndEvent) { /* e.Content */ })
|
||||
host.OnReasoningDelta(func(e kit.ReasoningDeltaEvent) { /* e.Delta */ })
|
||||
host.OnReasoningComplete(func(e kit.ReasoningCompleteEvent) {})
|
||||
host.OnToolExecutionStart(func(e kit.ToolExecutionStartEvent) { /* e.ToolCallID, e.ToolName, e.ToolKind, e.ToolArgs */ })
|
||||
host.OnToolExecutionEnd(func(e kit.ToolExecutionEndEvent) { /* e.ToolCallID, e.ToolName, e.ToolKind */ })
|
||||
host.OnToolCallContent(func(e kit.ToolCallContentEvent) { /* e.Content */ })
|
||||
host.OnStepUsage(func(e kit.StepUsageEvent) { /* e.InputTokens, e.OutputTokens, e.CacheReadTokens, e.CacheWriteTokens */ })
|
||||
host.OnCompaction(func(e kit.CompactionEvent) { /* e.Summary, e.OriginalTokens, e.CompactedTokens, ... */ })
|
||||
host.OnSteerConsumed(func(e kit.SteerConsumedEvent) { /* e.Count */ })
|
||||
```
|
||||
|
||||
> **Rename note:** `OnStreaming` has been renamed to `OnMessageUpdate`. The old `OnStreaming` name is kept as a deprecated alias for one release cycle.
|
||||
|
||||
### Generic subscriber (receives all events)
|
||||
|
||||
```go
|
||||
@@ -336,6 +392,16 @@ unsub := host.Subscribe(func(e kit.Event) {
|
||||
| `reasoning_delta` | `ReasoningDeltaEvent` | `Delta` |
|
||||
| `step_usage` | `StepUsageEvent` | `InputTokens`, `OutputTokens`, `CacheReadTokens`, `CacheWriteTokens` |
|
||||
| `steer_consumed` | `SteerConsumedEvent` | `Count` |
|
||||
| `step_start` | `StepStartEvent` | `StepNumber` |
|
||||
| `step_finish` | `StepFinishEvent` | `StepNumber`, `HasToolCalls`, `FinishReason`, `Usage` |
|
||||
| `text_start` | `TextStartEvent` | `ID` |
|
||||
| `text_end` | `TextEndEvent` | `ID` |
|
||||
| `reasoning_start` | `ReasoningStartEvent` | `ID` |
|
||||
| `warnings` | `WarningsEvent` | `Warnings` |
|
||||
| `source` | `SourceEvent` | `SourceType`, `ID`, `URL`, `Title` |
|
||||
| `stream_finish` | `StreamFinishEvent` | `Usage`, `FinishReason` |
|
||||
| `error` | `ErrorEvent` | `Error` |
|
||||
| `retry` | `RetryEvent` | `Attempt`, `Error` |
|
||||
| `password_prompt` | `PasswordPromptEvent` | `Prompt`, `ResponseCh` |
|
||||
|
||||
**Tool call streaming lifecycle**: `ToolCallStartEvent` → `ToolCallDeltaEvent` (repeated) → `ToolCallEndEvent` → `ToolCallEvent` → `ToolExecutionStartEvent` → `ToolOutputEvent` (optional, repeated) → `ToolExecutionEndEvent` → `ToolResultEvent`
|
||||
@@ -421,6 +487,20 @@ host.OnAfterTurn(kit.HookPriorityNormal, func(h kit.AfterTurnHook) {
|
||||
})
|
||||
```
|
||||
|
||||
### PrepareStep — intercept/replace messages before each LLM call
|
||||
|
||||
```go
|
||||
host.OnPrepareStep(kit.HookPriorityNormal, func(h kit.PrepareStepHook) *kit.PrepareStepResult {
|
||||
// h.StepNumber — which step in the current turn (1-based)
|
||||
// h.Messages — []kit.LLMMessage being sent to the LLM
|
||||
// Return nil to pass through unchanged, or replace messages:
|
||||
modified := filterSensitiveMessages(h.Messages)
|
||||
return &kit.PrepareStepResult{Messages: modified}
|
||||
})
|
||||
```
|
||||
|
||||
`PrepareStep` fires before every LLM API call within a turn (including tool-call loop iterations). Unlike `ContextPrepare` (which operates on the full context window once per turn), `PrepareStep` runs per-step and sees the messages that include the latest tool results.
|
||||
|
||||
### ContextPrepare — filter/inject context window
|
||||
|
||||
```go
|
||||
@@ -493,6 +573,8 @@ host, _ := kit.New(ctx, &kit.Options{
|
||||
|----------|-------------|
|
||||
| `kit.TextResult(content)` | Successful text result |
|
||||
| `kit.ErrorResult(content)` | Error result (LLM sees it as a tool error) |
|
||||
| `kit.ImageResult(content, data, mediaType)` | Image result with binary data (e.g. `"image/png"`) |
|
||||
| `kit.MediaResult(content, data, mediaType)` | Non-image media result (e.g. `"audio/mpeg"`) |
|
||||
|
||||
**ToolOutput fields** (for advanced use):
|
||||
|
||||
@@ -1095,6 +1177,8 @@ kit.LLMUsage // {InputTokens, OutputTokens, TotalTokens, ReasoningTokens,
|
||||
// CacheCreationTokens, CacheReadTokens}
|
||||
kit.LLMResponse // {Content, FinishReason, Usage}
|
||||
kit.LLMFilePart // {Filename, Data []byte, MediaType}
|
||||
kit.LLMToolCall // {ID, Name, Input string} — execution-layer tool call (for Tool.Run)
|
||||
kit.LLMToolResponse // {Type, Content, Data, MediaType, IsError, ...} — raw tool response
|
||||
|
||||
// Compaction types
|
||||
kit.CompactionResult, kit.CompactionOptions
|
||||
@@ -1168,7 +1252,7 @@ for {
|
||||
### Pattern: Streaming output to terminal
|
||||
|
||||
```go
|
||||
host.OnStreaming(func(e kit.MessageUpdateEvent) {
|
||||
host.OnMessageUpdate(func(e kit.MessageUpdateEvent) {
|
||||
fmt.Print(e.Chunk)
|
||||
})
|
||||
response, _ := host.Prompt(ctx, "Write a poem")
|
||||
|
||||
+54
-14
@@ -20,7 +20,7 @@ unsub2 := host.OnToolResult(func(event kit.ToolResultEvent) {
|
||||
})
|
||||
defer unsub2()
|
||||
|
||||
unsub3 := host.OnStreaming(func(event kit.MessageUpdateEvent) {
|
||||
unsub3 := host.OnMessageUpdate(func(event kit.MessageUpdateEvent) {
|
||||
fmt.Print(event.Chunk)
|
||||
})
|
||||
defer unsub3()
|
||||
@@ -116,6 +116,24 @@ host.OnAfterTurn(kit.HookPriorityNormal, func(h kit.AfterTurnHook) {
|
||||
})
|
||||
```
|
||||
|
||||
### PrepareStep — intercept messages between steps
|
||||
|
||||
The most powerful hook — fires between steps within a multi-step agent turn, after any steering messages are injected and before messages are sent to the LLM. Can replace the entire context window.
|
||||
|
||||
```go
|
||||
host.OnPrepareStep(kit.HookPriorityNormal, func(h kit.PrepareStepHook) *kit.PrepareStepResult {
|
||||
// h.StepNumber — zero-based step index within the turn
|
||||
// h.Messages — current context window (includes any steering)
|
||||
|
||||
// Example: transform tool results with images into user messages
|
||||
modified := transformImageToolResults(h.Messages)
|
||||
return &kit.PrepareStepResult{Messages: modified}
|
||||
// Return nil to pass through unchanged
|
||||
})
|
||||
```
|
||||
|
||||
Use cases: transforming tool results (e.g., image data for vision models), dynamic tool filtering per step, mid-turn context injection, custom stop conditions.
|
||||
|
||||
### Hook priorities
|
||||
|
||||
```go
|
||||
@@ -128,19 +146,41 @@ Lower values run first. First non-nil result wins.
|
||||
|
||||
## All event types
|
||||
|
||||
| Event | Description |
|
||||
|-------|-------------|
|
||||
| `ToolCallStartEvent` | LLM began generating tool call arguments (tool name known, args streaming) |
|
||||
| `ToolCallDeltaEvent` | Streamed JSON fragment of tool call arguments |
|
||||
| `ToolCallEndEvent` | Tool argument streaming complete, before execution begins |
|
||||
| `ToolCallEvent` | Tool call fully parsed and about to execute |
|
||||
| `ToolResultEvent` | Tool execution completed with result |
|
||||
| `ToolOutputEvent` | Streaming output chunk from tool (e.g., bash stdout/stderr) |
|
||||
| `MessageUpdateEvent` | Streaming text chunk from LLM |
|
||||
| `ResponseEvent` | Final response received |
|
||||
| `TurnStartEvent` | Agent turn started |
|
||||
| `TurnEndEvent` | Agent turn completed |
|
||||
| `PasswordPromptEvent` | Sudo command needs password (respond via `ResponseCh`) |
|
||||
| Event | Typed Subscriber | Description |
|
||||
|-------|-----------------|-------------|
|
||||
| `TurnStartEvent` | `OnTurnStart` | Agent turn started |
|
||||
| `TurnEndEvent` | `OnTurnEnd` | Agent turn completed |
|
||||
| `MessageStartEvent` | `OnMessageStart` | New assistant message begins |
|
||||
| `MessageUpdateEvent` | `OnMessageUpdate` | Streaming text chunk from LLM |
|
||||
| `MessageEndEvent` | `OnMessageEnd` | Assistant message complete |
|
||||
| `ToolCallStartEvent` | `OnToolCallStart` | LLM began generating tool call arguments |
|
||||
| `ToolCallDeltaEvent` | `OnToolCallDelta` | Streamed JSON fragment of tool call arguments |
|
||||
| `ToolCallEndEvent` | `OnToolCallEnd` | Tool argument streaming complete |
|
||||
| `ToolCallEvent` | `OnToolCall` | Tool call fully parsed, about to execute |
|
||||
| `ToolExecutionStartEvent` | `OnToolExecutionStart` | Tool begins executing |
|
||||
| `ToolExecutionEndEvent` | `OnToolExecutionEnd` | Tool finishes executing |
|
||||
| `ToolResultEvent` | `OnToolResult` | Tool execution completed with result |
|
||||
| `ToolCallContentEvent` | `OnToolCallContent` | Text content alongside tool calls |
|
||||
| `ToolOutputEvent` | `OnToolOutput` | Streaming output chunk from tool (e.g., bash) |
|
||||
| `ResponseEvent` | `OnResponse` | Final response received |
|
||||
| `ReasoningStartEvent` | `OnReasoningStart` | LLM begins reasoning/thinking |
|
||||
| `ReasoningDeltaEvent` | `OnReasoningDelta` | Streaming reasoning/thinking chunk |
|
||||
| `ReasoningCompleteEvent` | `OnReasoningComplete` | Reasoning/thinking finished |
|
||||
| `StepStartEvent` | `OnStepStart` | New LLM call begins within a turn |
|
||||
| `StepFinishEvent` | `OnStepFinish` | Step completes (with usage, finish reason, tool call info) |
|
||||
| `StepUsageEvent` | `OnStepUsage` | Per-step token usage |
|
||||
| `StreamFinishEvent` | `OnStreamFinish` | Per-step stream completes (with usage + finish reason) |
|
||||
| `TextStartEvent` | `OnTextStart` | LLM begins text content generation |
|
||||
| `TextEndEvent` | `OnTextEnd` | LLM finishes text content generation |
|
||||
| `WarningsEvent` | `OnWarnings` | LLM provider returned warnings |
|
||||
| `SourceEvent` | `OnSource` | LLM referenced a source (e.g., web search) |
|
||||
| `ErrorEvent` | `OnError` | Agent-level error during streaming |
|
||||
| `RetryEvent` | `OnRetry` | LLM request retried after transient error |
|
||||
| `CompactionEvent` | `OnCompaction` | Conversation compacted |
|
||||
| `SteerConsumedEvent` | `OnSteerConsumed` | Steering messages injected into turn |
|
||||
| `PasswordPromptEvent` | — | Sudo command needs password (respond via `ResponseCh`) |
|
||||
|
||||
> **Note:** `OnStreaming` is a deprecated alias for `OnMessageUpdate` and will be removed in a future release.
|
||||
|
||||
## Subagent event monitoring
|
||||
|
||||
|
||||
@@ -101,8 +101,10 @@ Return values:
|
||||
|--------|-------------|
|
||||
| `kit.TextResult(s)` | Successful text result |
|
||||
| `kit.ErrorResult(s)` | Error result (LLM sees it as a tool error) |
|
||||
| `kit.ImageResult(s, data, mediaType)` | Image result with binary data (e.g. `"image/png"`) |
|
||||
| `kit.MediaResult(s, data, mediaType)` | Non-image media result (e.g. `"audio/mpeg"`) |
|
||||
|
||||
For advanced use, return a `kit.ToolOutput` struct directly with `Data`, `MediaType`, and `Metadata` fields.
|
||||
Binary data (images, audio, etc.) in `ToolOutput.Data` is automatically forwarded to the LLM when `MediaType` is set. For advanced use, return a `kit.ToolOutput` struct directly with `Data`, `MediaType`, and `Metadata` fields.
|
||||
|
||||
Use `kit.NewParallelTool` for tools that are safe to run concurrently. Use `kit.ToolCallIDFromContext(ctx)` to retrieve the LLM-assigned call ID for logging or tracing.
|
||||
|
||||
@@ -141,7 +143,7 @@ host.OnToolResult(func(event kit.ToolResultEvent) {
|
||||
fmt.Println("Tool result:", event.Name)
|
||||
})
|
||||
|
||||
host.OnStreaming(func(event kit.MessageUpdateEvent) {
|
||||
host.OnMessageUpdate(func(event kit.MessageUpdateEvent) {
|
||||
fmt.Print(event.Chunk)
|
||||
})
|
||||
```
|
||||
|
||||
Reference in New Issue
Block a user