diff --git a/config/config.example.json b/config/config.example.json index 8db9dd3e9..e4a16de68 100644 --- a/config/config.example.json +++ b/config/config.example.json @@ -35,6 +35,12 @@ "model": "deepseek/deepseek-chat", "api_key": "sk-your-deepseek-key" }, + { + "model_name": "minimax", + "model": "minimax/MiniMax-M1", + "api_key": "your-minimax-api-key", + "stream": true + }, { "model_name": "loadbalanced-gpt4", "model": "openai/gpt-5.2", diff --git a/docs/plan-interview-improvements.md b/docs/plan-interview-improvements.md new file mode 100644 index 000000000..e083dfe48 --- /dev/null +++ b/docs/plan-interview-improvements.md @@ -0,0 +1,112 @@ +# /plan interview 改善検討 + +## 現象 + +`/plan ` でinterviewモードに入った後、AIが: + +1. MEMORY.mdに何も書かずに会話だけ続ける +2. interviewを無視して実装を始めようとする(exec, ファイル書き込み) +3. 結局フェーズ/ステップ/コマンドが書かれないまま executing に遷移する + +## 原因分析 + +### Layer 1: 技術的障壁(修正済み) + +| 問題 | 原因 | 修正 | +|---|---|---| +| XML tool callがパースされない | MiniMaxの開閉タグ不一致(`` vs ``) | regex + normalizeAlpha + 編集距離で fuzzy matching | +| tool名が不一致で実行失敗 | `readfile` vs `read_file` | `ToolRegistry.Get()` に normalizeAlpha フォールバック | +| interview許可リストも完全一致 | `isToolAllowedDuringInterview` が exact match | normalizeAlpha で比較 | + +### Layer 2: AI行動の問題(未対応) + +技術的障壁を除去しても、モデル(特に小規模モデル)の命令追従に起因する問題が残る: + +- **「会話しながらファイル編集」が複合タスクとして難しい** — interview中にMEMORY.mdを更新する行為は、会話とファイル操作の並行処理。小さいモデルにはハードルが高い +- **AIがワークフローを無視して実装に走る** — interview指示よりも「ユーザーの要求を直接解決しよう」というバイアスが強い +- **tool callイテレーション中に目的を忘れる** — read_fileでファイルを読み始めると、そのまま実装に入ろうとする + +## 検討した案と判断 + +### 案: ユーザー発言の自動追記(却下) + +MEMORY.mdにユーザー発言を自動追記 → ちゃんとしたplanにはならない。生データの蓄積であって構造化された計画ではない。 + +### 案: 専用tool `save_context`(却下) + +tool仕様を変えても、AIがtoolを適切に呼ばない根本問題は解決しない。 + +### 案: 別パスでplan生成LLMコール(却下) + +AIが「情報が揃った」と判断するトリガーの設計が難しい。キーワード検出は不安定。 + +### 案: `/plan draft` コマンド(却下) + +ユーザーがトリガーできても、その前にAIがワークフローを無視して実装を始める問題は残る。また状態とコマンドが増えてユーザーが混乱する。 + +### 案: 毎ターン固定メッセージ注入(却下) + +探索的なマルチターン会話で誤爆する。AIがファイルを読んだりリサーチしている途中のターンで「edit_fileしろ」は邪魔。 + +## 採用方針: tool callイテレーション内リマインド + +### 既知の知見 + +通常の開発モード(executing)で、tool callが反復される中でユーザー指示が忘れられる問題に対し、リマインド注入で自律開発がスムーズに進むようになった実績がある。同じパターンをinterview中にも適用する。 + +### 設計 + +**1. interviewフェーズのtool制限(実装済み)** + +``` +許可: read_file, list_dir, web_search, web_fetch +許可: edit_file / write_file / append_file(MEMORY.mdのみ) +ブロック: exec, その他write系 +``` + +AIが実装に走ろうとしても物理的にできない。 + +**2. tool callイテレーション内でリマインド注入(未実装)** + +`runLLMIteration` 内で、tool結果をLLMに返す直前(= 次のLLMコールの直前)にリマインドを差し込む。 + +```go +// tool結果メッセージの後、次のLLMコール前 +if isPlanPreExecution(agent.ContextBuilder.GetPlanStatus()) { + messages = append(messages, providers.Message{ + Role: "user", + Content: "[System] You are interviewing. Ask questions and save findings " + + "to ## Context in memory/MEMORY.md. " + + "When ready, write ## Phase and ## Commands sections.", + }) +} +``` + +- **注入タイミング**: tool callループ内のみ。ユーザーとの会話ターンには入れない +- **注入条件**: interviewing または review 状態の時 +- **内容**: 固定。短く、具体的に何をすべきか指示 + +**3. 状態遷移は既存のまま** + +``` +/plan → interviewing(AIが質問、read系+MEMORY.md書き込み許可) + → AIがStatus:executingに変更しようとする + → システムがreviewに横取り(phases > 0 の場合) + → ユーザーにplan表示 +/plan start → executing(全toolアンロック) +``` + +新しい状態・新しいコマンドなし。 + +## 実装タスク + +- [ ] `runLLMIteration` 内のtool callループにリマインド注入を追加 +- [ ] リマインド内容をステータスごとに分岐(interviewing / review / executing) +- [ ] 既存のstaleness nudge(2ターン無更新で警告)との統合・整理 +- [ ] テスト追加 + +## 未解決の懸念 + +- **リマインドの効果がモデル依存**: 大きいモデルには効くが、小さいモデルでは無視される可能性 +- **リマインドの頻度**: 毎イテレーション注入でトークン消費が増える(ただし1行程度なので軽微) +- **interview→plan書き込みのタイミング**: AIが「もう十分」と判断する基準はモデル任せ。staleness nudgeが補助するが確実ではない diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 7212f9883..673f91aa1 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -26,6 +26,7 @@ import ( "github.com/sipeed/picoclaw/pkg/constants" "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/providers" + "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" "github.com/sipeed/picoclaw/pkg/routing" "github.com/sipeed/picoclaw/pkg/session" "github.com/sipeed/picoclaw/pkg/skills" @@ -258,6 +259,14 @@ func (al *AgentLoop) Run(ctx context.Context) error { Content: "via MiniApp: " + msg.Content, SkipPlaceholder: true, }) + // Create a placeholder AFTER the echo so status updates appear below it. + if al.channelManager != nil { + if ch, ok := al.channelManager.GetChannel(msg.Channel); ok { + if tc, ok := ch.(*channels.TelegramChannel); ok { + tc.CreatePlaceholder(ctx, msg.ChatID) + } + } + } } // Fast path: handle slash commands immediately without blocking the LLM worker. @@ -824,10 +833,10 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt // 5a. Auto-advance plan phases after LLM iteration postStatus := agent.ContextBuilder.GetPlanStatus() - if agent.ContextBuilder.HasActivePlan() && postStatus == "executing" { - // Intercept: if AI changed status to executing without user approval - // (from interviewing or review), validate and set to "review". - if preStatus == "interviewing" || preStatus == "review" { + if agent.ContextBuilder.HasActivePlan() && (postStatus == "executing" || postStatus == "review") { + // Intercept: if AI changed status to executing or review without user approval + // (from interviewing or review), validate and hold at "review". + if preStatus == "interviewing" || (preStatus == "review" && postStatus == "executing") { if err := agent.ContextBuilder.ValidatePlanStructure(); err != nil { _ = agent.ContextBuilder.SetPlanStatus("interviewing") logger.WarnCF("agent", "Reverted plan to interviewing: "+err.Error(), @@ -847,7 +856,7 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt }) } } - } else if agent.ContextBuilder.GetTotalPhases() == 0 { + } else if postStatus == "executing" && agent.ContextBuilder.GetTotalPhases() == 0 { // Safeguard: executing but no phases (shouldn't happen, but be safe). _ = agent.ContextBuilder.SetPlanStatus("interviewing") logger.WarnCF("agent", "Reverted plan to interviewing: no phases defined", @@ -1237,9 +1246,10 @@ func compressRepeats(s string) string { // Display layout constants. const ( - displayPastEntries = 4 // number of compact 1-line past entries - displayErrorLines = 5 // content lines inside the error code block - statusSeparator = "\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\n" + displayPastEntries = 4 // number of compact 1-line past entries + displayErrorLines = 5 // content lines inside the error code block + statusSeparator = "\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\n" + streamingDisplayLines = 17 // line count matching buildRichStatus output ) // buildRichStatus builds a fixed-height terminal-like status display. @@ -1377,6 +1387,101 @@ func buildRichStatus(task *activeTask, isBackground bool, workspace string) stri } // runLLMIteration executes the LLM call loop with tool handling. +// consumeStreamWithRepetitionDetection reads StreamEvents from ch, accumulates +// content and tool calls, and runs repetition detection every checkInterval runes. +// If repetition is detected, cancelFn is called to abort the HTTP request and +// the function returns the partial response with detected=true. +func consumeStreamWithRepetitionDetection( + ch <-chan protocoltypes.StreamEvent, + cancelFn context.CancelFunc, + checkInterval int, + onChunk func(accumulated string), +) (*providers.LLMResponse, bool, error) { + var content strings.Builder + var toolCalls []streamToolCallAcc + var finishReason string + var usage *providers.UsageInfo + runesSinceLastCheck := 0 + + for ev := range ch { + if ev.Err != nil { + return nil, false, ev.Err + } + if ev.ContentDelta != "" { + content.WriteString(ev.ContentDelta) + runesSinceLastCheck += utf8.RuneCountInString(ev.ContentDelta) + if onChunk != nil { + onChunk(content.String()) + } + } + if ev.FinishReason != "" { + finishReason = ev.FinishReason + } + if ev.Usage != nil { + usage = ev.Usage + } + for _, tc := range ev.ToolCallDeltas { + for len(toolCalls) <= tc.Index { + toolCalls = append(toolCalls, streamToolCallAcc{}) + } + if tc.ID != "" { + toolCalls[tc.Index].id = tc.ID + } + if tc.Name != "" { + toolCalls[tc.Index].name = tc.Name + } + toolCalls[tc.Index].args.WriteString(tc.ArgumentsDelta) + } + + // Run repetition detection periodically on accumulated content. + if runesSinceLastCheck >= checkInterval && content.Len() > 2000 { + runesSinceLastCheck = 0 + if utils.DetectRepetitionLoop(content.String()) { + cancelFn() + // Drain remaining events so the producer goroutine can exit. + for range ch { + } + resp := buildAccumulatedResponse(content.String(), toolCalls, finishReason, usage) + return resp, true, nil + } + } + } + + resp := buildAccumulatedResponse(content.String(), toolCalls, finishReason, usage) + return resp, false, nil +} + +// streamToolCallAcc accumulates streamed tool call fragments. +type streamToolCallAcc struct { + id string + name string + args strings.Builder +} + +// buildAccumulatedResponse constructs an LLMResponse from accumulated stream data. +func buildAccumulatedResponse(content string, toolCalls []streamToolCallAcc, finishReason string, usage *providers.UsageInfo) *providers.LLMResponse { + resp := &providers.LLMResponse{ + Content: content, + FinishReason: finishReason, + Usage: usage, + } + for _, tc := range toolCalls { + arguments := make(map[string]any) + argStr := tc.args.String() + if argStr != "" { + if err := json.Unmarshal([]byte(argStr), &arguments); err != nil { + arguments["raw"] = argStr + } + } + resp.ToolCalls = append(resp.ToolCalls, providers.ToolCall{ + ID: tc.id, + Name: tc.name, + Arguments: arguments, + }) + } + return resp +} + func (al *AgentLoop) runLLMIteration( ctx context.Context, agent *AgentInstance, @@ -1459,15 +1564,59 @@ func (al *AgentLoop) runLLMIteration( var response *providers.LLMResponse var err error + // Build onChunk callback for streaming preview. + // When sending responses to a real (non-internal) channel, publish + // throttled status updates so the user sees LLM output in real time. + var onChunk func(string) + if !constants.IsInternalChannel(opts.Channel) { + lastPublish := time.Time{} + onChunk = func(accumulated string) { + if time.Since(lastPublish) < 500*time.Millisecond { + return + } + lastPublish = time.Now() + display := utils.TailPad(accumulated, streamingDisplayLines, maxEntryLineWidth) + al.bus.PublishOutbound(bus.OutboundMessage{ + Channel: opts.Channel, + ChatID: opts.ChatID, + Content: display + " \u2589", + IsStatus: true, + }) + } + } + + // doCall invokes a single LLM provider, using streaming with + // early repetition detection when the provider supports it. + opts_ := map[string]any{ + "max_tokens": agent.MaxTokens, + "temperature": agent.Temperature, + } + doCall := func(ctx context.Context, p providers.LLMProvider, model string) (*providers.LLMResponse, error) { + if sp, ok := p.(providers.StreamingProvider); ok && sp.CanStream() { + streamCtx, streamCancel := context.WithCancel(ctx) + defer streamCancel() + ch, err := sp.ChatStream(streamCtx, messages, providerToolDefs, model, opts_) + if err != nil { + return nil, err + } + resp, repetition, err := consumeStreamWithRepetitionDetection(ch, streamCancel, 1000, onChunk) + if err != nil { + return nil, err + } + if repetition { + resp.FinishReason = "repetition_detected" + } + return resp, nil + } + return p.Chat(ctx, messages, providerToolDefs, model, opts_) + } + callLLM := func() (*providers.LLMResponse, error) { if len(agent.Candidates) > 1 && al.fallback != nil { fbResult, fbErr := al.fallback.Execute(ctx, agent.Candidates, func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) { p := al.resolveProvider(provider, model, agent.Provider) - return p.Chat(ctx, messages, providerToolDefs, model, map[string]interface{}{ - "max_tokens": agent.MaxTokens, - "temperature": agent.Temperature, - }) + return doCall(ctx, p, model) }, ) if fbErr != nil { @@ -1480,10 +1629,7 @@ func (al *AgentLoop) runLLMIteration( } return fbResult.Response, nil } - return agent.Provider.Chat(ctx, messages, providerToolDefs, agent.Model, map[string]any{ - "max_tokens": agent.MaxTokens, - "temperature": agent.Temperature, - }) + return doCall(ctx, agent.Provider, agent.Model) } // Retry loop for context/token errors @@ -1545,6 +1691,47 @@ func (al *AgentLoop) runLLMIteration( ) } + // Detect repetition loop on raw text (before stripping think + // blocks so loops inside are caught). Skip when the + // provider already returned native tool calls. + // Streaming providers may have already flagged repetition via + // FinishReason="repetition_detected" — honour that too. + if response.FinishReason == "repetition_detected" || + (len(response.ToolCalls) == 0 && utils.DetectRepetitionLoop(response.Content)) { + logger.WarnCF("agent", "Repetition loop detected in LLM response, retrying", + map[string]any{ + "agent_id": agent.ID, + "iteration": iteration, + "finish_reason": response.FinishReason, + "content_length": len(response.Content), + }) + + // Retry once: inject nudge message and re-call + savedMsgs := messages + messages = append(append([]providers.Message(nil), messages...), + providers.Message{ + Role: "user", + Content: "[System] Your previous response contained degenerate repetition and was discarded. Please respond normally without repeating yourself.", + }) + response, err = callLLM() + messages = savedMsgs // restore original messages + + if err != nil { + return "", iteration, fmt.Errorf("LLM retry after repetition failed: %w", err) + } + + // Re-check on raw text; if still repeating give up + if utils.DetectRepetitionLoop(response.Content) { + logger.ErrorCF("agent", "Repetition persists after retry, returning empty", + map[string]any{"agent_id": agent.ID}) + response.Content = "" + } + } + + // Strip think blocks before extracting XML tool calls so + // extraction operates on clean content. + response.Content = utils.StripThinkBlocks(response.Content) + // Recover XML tool calls emitted as plain text by some providers. if len(response.ToolCalls) == 0 { if xmlCalls := providers.ExtractXMLToolCalls(response.Content); len(xmlCalls) > 0 { diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index fe66b6d94..0e42b98ad 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -12,6 +12,7 @@ import ( "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/providers" + "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" "github.com/sipeed/picoclaw/pkg/tools" ) @@ -2116,3 +2117,228 @@ func (m *nudgeCaptureMockProvider) Chat( func (m *nudgeCaptureMockProvider) GetDefaultModel() string { return "mock-nudge-model" } + +// --- consumeStreamWithRepetitionDetection tests --- + +func TestConsumeStream_NormalCompletion(t *testing.T) { + ch := make(chan protocoltypes.StreamEvent, 8) + go func() { + ch <- protocoltypes.StreamEvent{ContentDelta: "Hello "} + ch <- protocoltypes.StreamEvent{ContentDelta: "world!"} + ch <- protocoltypes.StreamEvent{ + FinishReason: "stop", + Usage: &providers.UsageInfo{PromptTokens: 5, CompletionTokens: 2, TotalTokens: 7}, + } + close(ch) + }() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + resp, detected, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000, nil) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if detected { + t.Fatal("expected detected=false for normal content") + } + if resp.Content != "Hello world!" { + t.Errorf("Content = %q, want %q", resp.Content, "Hello world!") + } + if resp.FinishReason != "stop" { + t.Errorf("FinishReason = %q, want %q", resp.FinishReason, "stop") + } + if resp.Usage == nil || resp.Usage.TotalTokens != 7 { + t.Errorf("Usage.TotalTokens = %v, want 7", resp.Usage) + } + _ = ctx // keep linter happy +} + +func TestConsumeStream_DetectsRepetition(t *testing.T) { + ch := make(chan protocoltypes.StreamEvent, 64) + cancelCalled := false + + ctx, cancel := context.WithCancel(context.Background()) + wrappedCancel := func() { + cancelCalled = true + cancel() + } + + // Send enough repetitive content to trigger detection. + // The pattern "abcdefghij" repeated many times will have very low n-gram uniqueness. + repeatedChunk := strings.Repeat("abcdefghij", 50) // 500 chars per chunk + go func() { + // Send 6 chunks of repetitive content = 3000 chars total, + // each with 500 runes. The check triggers after every 1000 runes + // when content > 2000 chars. + for i := 0; i < 6; i++ { + ch <- protocoltypes.StreamEvent{ContentDelta: repeatedChunk} + } + // Send more data that should be ignored after detection. + for i := 0; i < 10; i++ { + ch <- protocoltypes.StreamEvent{ContentDelta: "more data"} + } + close(ch) + }() + + resp, detected, err := consumeStreamWithRepetitionDetection(ch, wrappedCancel, 1000, nil) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !detected { + t.Fatal("expected repetition detection to trigger") + } + if !cancelCalled { + t.Error("expected cancelFn to be called") + } + // The response should be shorter than the full 3000+ chars + // because detection triggers early. + if len(resp.Content) >= 3000+10*len("more data") { + t.Errorf("Content length = %d, expected less than full output", len(resp.Content)) + } + _ = ctx +} + +func TestConsumeStream_ToolCallAccumulation(t *testing.T) { + ch := make(chan protocoltypes.StreamEvent, 8) + go func() { + ch <- protocoltypes.StreamEvent{ + ToolCallDeltas: []protocoltypes.StreamToolCallDelta{ + {Index: 0, ID: "call_1", Name: "test_fn", ArgumentsDelta: `{"ke`}, + }, + } + ch <- protocoltypes.StreamEvent{ + ToolCallDeltas: []protocoltypes.StreamToolCallDelta{ + {Index: 0, ArgumentsDelta: `y":"val"}`}, + }, + } + ch <- protocoltypes.StreamEvent{FinishReason: "tool_calls"} + close(ch) + }() + + _, cancel := context.WithCancel(context.Background()) + defer cancel() + + resp, detected, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000, nil) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if detected { + t.Fatal("expected no repetition detection for tool calls") + } + if len(resp.ToolCalls) != 1 { + t.Fatalf("len(ToolCalls) = %d, want 1", len(resp.ToolCalls)) + } + if resp.ToolCalls[0].Name != "test_fn" { + t.Errorf("ToolCalls[0].Name = %q, want %q", resp.ToolCalls[0].Name, "test_fn") + } + if resp.ToolCalls[0].Arguments["key"] != "val" { + t.Errorf("ToolCalls[0].Arguments[key] = %v, want %q", resp.ToolCalls[0].Arguments["key"], "val") + } +} + +func TestConsumeStream_StreamError(t *testing.T) { + ch := make(chan protocoltypes.StreamEvent, 4) + go func() { + ch <- protocoltypes.StreamEvent{ContentDelta: "partial"} + ch <- protocoltypes.StreamEvent{Err: fmt.Errorf("read error")} + close(ch) + }() + + _, cancel := context.WithCancel(context.Background()) + defer cancel() + + _, _, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000, nil) + if err == nil { + t.Fatal("expected error, got nil") + } + if !strings.Contains(err.Error(), "read error") { + t.Errorf("error = %q, want to contain %q", err.Error(), "read error") + } +} + +func TestConsumeStream_OnChunkCallback(t *testing.T) { + ch := make(chan protocoltypes.StreamEvent, 8) + go func() { + ch <- protocoltypes.StreamEvent{ContentDelta: "Hello "} + ch <- protocoltypes.StreamEvent{ContentDelta: "world"} + ch <- protocoltypes.StreamEvent{ContentDelta: "!"} + ch <- protocoltypes.StreamEvent{FinishReason: "stop"} + close(ch) + }() + + _, cancel := context.WithCancel(context.Background()) + defer cancel() + + var chunks []string + onChunk := func(accumulated string) { + chunks = append(chunks, accumulated) + } + + resp, detected, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000, onChunk) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if detected { + t.Fatal("expected detected=false") + } + if resp.Content != "Hello world!" { + t.Errorf("Content = %q, want %q", resp.Content, "Hello world!") + } + // onChunk should be called once per content delta (3 times) + if len(chunks) != 3 { + t.Fatalf("onChunk called %d times, want 3", len(chunks)) + } + if chunks[0] != "Hello " { + t.Errorf("chunks[0] = %q, want %q", chunks[0], "Hello ") + } + if chunks[1] != "Hello world" { + t.Errorf("chunks[1] = %q, want %q", chunks[1], "Hello world") + } + if chunks[2] != "Hello world!" { + t.Errorf("chunks[2] = %q, want %q", chunks[2], "Hello world!") + } +} + +func TestConsumeStream_OnChunkWithRepetitionDetection(t *testing.T) { + ch := make(chan protocoltypes.StreamEvent, 64) + cancelCalled := false + + ctx, cancel := context.WithCancel(context.Background()) + wrappedCancel := func() { + cancelCalled = true + cancel() + } + + repeatedChunk := strings.Repeat("abcdefghij", 50) // 500 chars per chunk + go func() { + for i := 0; i < 6; i++ { + ch <- protocoltypes.StreamEvent{ContentDelta: repeatedChunk} + } + for i := 0; i < 10; i++ { + ch <- protocoltypes.StreamEvent{ContentDelta: "more data"} + } + close(ch) + }() + + var chunkCount int + onChunk := func(accumulated string) { + chunkCount++ + } + + _, detected, err := consumeStreamWithRepetitionDetection(ch, wrappedCancel, 1000, onChunk) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !detected { + t.Fatal("expected repetition detection to trigger") + } + if !cancelCalled { + t.Error("expected cancelFn to be called") + } + // onChunk should have been called at least once before detection + if chunkCount == 0 { + t.Error("expected onChunk to be called at least once") + } + _ = ctx +} diff --git a/pkg/agent/memory.go b/pkg/agent/memory.go index 05ad547f7..4cb825898 100644 --- a/pkg/agent/memory.go +++ b/pkg/agent/memory.go @@ -489,7 +489,7 @@ func (ms *MemoryStore) GetInterviewContext() string { sb.WriteString("- When you have enough information, use edit_file to add ## Phase, ## Commands, and ## Context sections BELOW the header block.\n") sb.WriteString("- Each step MUST use checkbox syntax: `- [ ] description`. The system parses checkboxes to track progress.\n") sb.WriteString("- Organize into 2-5 phases with 3-5 steps each.\n") - sb.WriteString("- After writing Phases, change `> Status: interviewing` to `> Status: executing` via edit_file.\n") + sb.WriteString("- After writing Phases, change `> Status: interviewing` to `> Status: review` via edit_file. The user must approve with /plan start before execution begins.\n") sb.WriteString("\n### Target Format (MANDATORY — system parses this exact structure)\n") sb.WriteString("\n") sb.WriteString("# Active Plan\n") diff --git a/pkg/channels/telegram.go b/pkg/channels/telegram.go index 006daf11d..f1a0f5266 100644 --- a/pkg/channels/telegram.go +++ b/pkg/channels/telegram.go @@ -51,8 +51,6 @@ const telegramMaxMessageChars = 3900 const markdownTableMaxWidth = 42 const markdownTableMinColWidth = 6 -var thinkBlockPattern = regexp.MustCompile(`(?is).*?`) - func NewTelegramChannel(cfg *config.Config, bus *bus.MessageBus) (*TelegramChannel, error) { var opts []telego.BotOption telegramCfg := cfg.Channels.Telegram @@ -196,6 +194,35 @@ func (c *TelegramChannel) Stop(ctx context.Context) error { return nil } +// CreatePlaceholder sends a "Thinking..." placeholder for the given chatID +// so that subsequent IsStatus messages can update it via EditStatus. +func (c *TelegramChannel) CreatePlaceholder(ctx context.Context, chatID string) error { + if !c.IsRunning() { + return nil + } + cid, err := parseChatID(chatID) + if err != nil { + return err + } + + // Stop any previous thinking animation + if prevStop, ok := c.stopThinking.Load(chatID); ok { + if cf, ok := prevStop.(*thinkingCancel); ok && cf != nil { + cf.Cancel() + } + } + + _, thinkCancel := context.WithTimeout(ctx, 5*time.Minute) + c.stopThinking.Store(chatID, &thinkingCancel{fn: thinkCancel}) + + pMsg, err := c.bot.SendMessage(ctx, tu.Message(tu.ID(cid), "Thinking... 💭")) + if err != nil { + return err + } + c.placeholders.Store(chatID, pMsg.MessageID) + return nil +} + func (c *TelegramChannel) EditStatus(ctx context.Context, msg bus.OutboundMessage) error { if !c.IsRunning() { return nil @@ -326,8 +353,7 @@ func (c *TelegramChannel) Send(ctx context.Context, msg bus.OutboundMessage) err } func sanitizeTelegramOutgoingContent(content string) string { - cleaned := thinkBlockPattern.ReplaceAllString(content, "") - cleaned = strings.TrimSpace(cleaned) + cleaned := strings.TrimSpace(content) if cleaned == "" { return "(empty response)" } diff --git a/pkg/channels/telegram_test.go b/pkg/channels/telegram_test.go index 39b59d0e3..3f3ccaca8 100644 --- a/pkg/channels/telegram_test.go +++ b/pkg/channels/telegram_test.go @@ -5,8 +5,8 @@ import ( "testing" ) -func TestSanitizeTelegramOutgoingContent_RemovesThinkBlock(t *testing.T) { - in := "\nsecret reasoning\n\n\nユーザー向け本文" +func TestSanitizeTelegramOutgoingContent_PlainText(t *testing.T) { + in := " ユーザー向け本文 " got := sanitizeTelegramOutgoingContent(in) want := "ユーザー向け本文" if got != want { @@ -14,9 +14,16 @@ func TestSanitizeTelegramOutgoingContent_RemovesThinkBlock(t *testing.T) { } } -func TestSanitizeTelegramOutgoingContent_EmptyAfterThink(t *testing.T) { - in := "only reasoning" - got := sanitizeTelegramOutgoingContent(in) +func TestSanitizeTelegramOutgoingContent_Empty(t *testing.T) { + got := sanitizeTelegramOutgoingContent("") + want := "(empty response)" + if got != want { + t.Fatalf("sanitizeTelegramOutgoingContent() = %q, want %q", got, want) + } +} + +func TestSanitizeTelegramOutgoingContent_WhitespaceOnly(t *testing.T) { + got := sanitizeTelegramOutgoingContent(" \n\t ") want := "(empty response)" if got != want { t.Fatalf("sanitizeTelegramOutgoingContent() = %q, want %q", got, want) diff --git a/pkg/config/config.go b/pkg/config/config.go index 6ce0b6ad6..aec93f663 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -397,6 +397,7 @@ type ModelConfig struct { // Optional optimizations RPM int `json:"rpm,omitempty"` // Requests per minute limit MaxTokensField string `json:"max_tokens_field,omitempty"` // Field name for max tokens (e.g., "max_completion_tokens") + Stream *bool `json:"stream,omitempty"` // Use SSE streaming (default: protocol-dependent) } // Validate checks if the ModelConfig has all required fields. diff --git a/pkg/providers/factory_provider.go b/pkg/providers/factory_provider.go index 74fe8a36c..a0e44aca7 100644 --- a/pkg/providers/factory_provider.go +++ b/pkg/providers/factory_provider.go @@ -10,6 +10,7 @@ import ( "strings" "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/providers/openai_compat" ) // createClaudeAuthProvider creates a Claude provider using OAuth credentials from auth store. @@ -84,7 +85,25 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err if apiBase == "" { apiBase = getDefaultAPIBase(protocol) } - return NewHTTPProviderWithMaxTokensField(cfg.APIKey, apiBase, cfg.Proxy, cfg.MaxTokensField), modelID, nil + return NewHTTPProviderWithOptions(cfg.APIKey, apiBase, cfg.Proxy, openai_compat.Options{ + MaxTokensField: cfg.MaxTokensField, + Stream: boolDefault(cfg.Stream, false), + }), modelID, nil + + case "minimax": + // MiniMax uses a non-standard endpoint path and defaults to SSE streaming. + if cfg.APIKey == "" && cfg.APIBase == "" { + return nil, "", fmt.Errorf("api_key or api_base is required for minimax protocol") + } + apiBase := cfg.APIBase + if apiBase == "" { + apiBase = getDefaultAPIBase(protocol) + } + return NewHTTPProviderWithOptions(cfg.APIKey, apiBase, cfg.Proxy, openai_compat.Options{ + EndpointPath: "/text/chatcompletion_v2", + MaxTokensField: cfg.MaxTokensField, + Stream: boolDefault(cfg.Stream, true), + }), modelID, nil case "openrouter", "groq", "zhipu", "gemini", "nvidia", "ollama", "moonshot", "shengsuanyun", "deepseek", "cerebras", @@ -97,7 +116,10 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err if apiBase == "" { apiBase = getDefaultAPIBase(protocol) } - return NewHTTPProviderWithMaxTokensField(cfg.APIKey, apiBase, cfg.Proxy, cfg.MaxTokensField), modelID, nil + return NewHTTPProviderWithOptions(cfg.APIKey, apiBase, cfg.Proxy, openai_compat.Options{ + MaxTokensField: cfg.MaxTokensField, + Stream: boolDefault(cfg.Stream, false), + }), modelID, nil case "anthropic": if cfg.AuthMethod == "oauth" || cfg.AuthMethod == "token" { @@ -186,7 +208,17 @@ func getDefaultAPIBase(protocol string) string { return "https://dashscope.aliyuncs.com/compatible-mode/v1" case "vllm": return "http://localhost:8000/v1" + case "minimax": + return "https://api.minimax.io/v1" default: return "" } } + +// boolDefault dereferences a *bool, returning def when nil. +func boolDefault(p *bool, def bool) bool { + if p != nil { + return *p + } + return def +} diff --git a/pkg/providers/http_provider.go b/pkg/providers/http_provider.go index f669d74cb..7572a3837 100644 --- a/pkg/providers/http_provider.go +++ b/pkg/providers/http_provider.go @@ -28,6 +28,12 @@ func NewHTTPProviderWithMaxTokensField(apiKey, apiBase, proxy, maxTokensField st } } +func NewHTTPProviderWithOptions(apiKey, apiBase, proxy string, opts openai_compat.Options) *HTTPProvider { + return &HTTPProvider{ + delegate: openai_compat.NewProviderWithOptions(apiKey, apiBase, proxy, opts), + } +} + func (p *HTTPProvider) Chat(ctx context.Context, messages []Message, tools []ToolDefinition, model string, options map[string]interface{}) (*LLMResponse, error) { resp, err := p.delegate.Chat(ctx, messages, tools, model, options) if err != nil { @@ -48,3 +54,19 @@ func (p *HTTPProvider) Chat(ctx context.Context, messages []Message, tools []Too func (p *HTTPProvider) GetDefaultModel() string { return "" } + +// CanStream returns true when the underlying provider uses SSE streaming. +func (p *HTTPProvider) CanStream() bool { + return p.delegate.CanStream() +} + +// ChatStream opens an SSE stream and returns a channel of StreamEvent. +func (p *HTTPProvider) ChatStream( + ctx context.Context, + messages []Message, + tools []ToolDefinition, + model string, + options map[string]any, +) (<-chan StreamEvent, error) { + return p.delegate.ChatStream(ctx, messages, tools, model, options) +} diff --git a/pkg/providers/openai_compat/provider.go b/pkg/providers/openai_compat/provider.go index b8528953a..2c49e1044 100644 --- a/pkg/providers/openai_compat/provider.go +++ b/pkg/providers/openai_compat/provider.go @@ -1,6 +1,7 @@ package openai_compat import ( + "bufio" "bytes" "context" "encoding/json" @@ -30,17 +31,36 @@ type ( type Provider struct { apiKey string apiBase string + endpointPath string // API path appended to apiBase (default: "/chat/completions") maxTokensField string // Field name for max tokens (e.g., "max_completion_tokens" for o1/glm models) + stream bool // Use SSE streaming internally (accumulates into a single LLMResponse) httpClient *http.Client } +// Options configures optional behaviour for the provider. +type Options struct { + EndpointPath string // API path appended to apiBase (default: "/chat/completions") + MaxTokensField string // Field name for max tokens parameter + Stream bool // Use SSE streaming internally +} + func NewProvider(apiKey, apiBase, proxy string) *Provider { return NewProviderWithMaxTokensField(apiKey, apiBase, proxy, "") } func NewProviderWithMaxTokensField(apiKey, apiBase, proxy, maxTokensField string) *Provider { + return NewProviderWithOptions(apiKey, apiBase, proxy, Options{ + MaxTokensField: maxTokensField, + }) +} + +func NewProviderWithOptions(apiKey, apiBase, proxy string, opts Options) *Provider { + timeout := 120 * time.Second + if opts.Stream { + timeout = 5 * time.Minute + } client := &http.Client{ - Timeout: 120 * time.Second, + Timeout: timeout, } if proxy != "" { @@ -54,21 +74,33 @@ func NewProviderWithMaxTokensField(apiKey, apiBase, proxy, maxTokensField string } } + endpointPath := opts.EndpointPath + if endpointPath == "" { + endpointPath = "/chat/completions" + } + return &Provider{ apiKey: apiKey, apiBase: strings.TrimRight(apiBase, "/"), - maxTokensField: maxTokensField, + endpointPath: endpointPath, + maxTokensField: opts.MaxTokensField, + stream: opts.Stream, httpClient: client, } } -func (p *Provider) Chat( +// streamBufferSize is the channel buffer size for ChatStream events. +const streamBufferSize = 32 + +// buildHTTPRequest constructs a ready-to-send *http.Request for the chat API. +func (p *Provider) buildHTTPRequest( ctx context.Context, messages []Message, tools []ToolDefinition, model string, options map[string]any, -) (*LLMResponse, error) { + stream bool, +) (*http.Request, error) { if p.apiBase == "" { return nil, fmt.Errorf("API base not configured") } @@ -111,12 +143,16 @@ func (p *Provider) Chat( } } + if stream { + requestBody["stream"] = true + } + jsonData, err := json.Marshal(requestBody) if err != nil { return nil, fmt.Errorf("failed to marshal request: %w", err) } - req, err := http.NewRequestWithContext(ctx, "POST", p.apiBase+"/chat/completions", bytes.NewReader(jsonData)) + req, err := http.NewRequestWithContext(ctx, "POST", p.apiBase+p.endpointPath, bytes.NewReader(jsonData)) if err != nil { return nil, fmt.Errorf("failed to create request: %w", err) } @@ -126,22 +162,216 @@ func (p *Provider) Chat( req.Header.Set("Authorization", "Bearer "+p.apiKey) } + return req, nil +} + +func (p *Provider) Chat( + ctx context.Context, + messages []Message, + tools []ToolDefinition, + model string, + options map[string]any, +) (*LLMResponse, error) { + // When streaming is enabled, delegate to ChatStream + AccumulateStream + // so that the SSE→channel path is always exercised. + if p.stream { + ch, err := p.ChatStream(ctx, messages, tools, model, options) + if err != nil { + return nil, err + } + return AccumulateStream(ch) + } + + req, err := p.buildHTTPRequest(ctx, messages, tools, model, options, false) + if err != nil { + return nil, err + } + resp, err := p.httpClient.Do(req) if err != nil { return nil, fmt.Errorf("failed to send request: %w", err) } defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("API request failed:\n Status: %d\n Body: %s", resp.StatusCode, string(body)) + } + body, err := io.ReadAll(resp.Body) if err != nil { return nil, fmt.Errorf("failed to read response: %w", err) } + return parseResponse(body) +} + +// CanStream returns true when this provider is configured for SSE streaming. +func (p *Provider) CanStream() bool { + return p.stream +} + +// ChatStream opens an SSE connection and returns a channel of StreamEvent. +// The channel is closed when the stream ends or an error occurs. +// Cancelling ctx will abort the HTTP request and close the channel. +func (p *Provider) ChatStream( + ctx context.Context, + messages []Message, + tools []ToolDefinition, + model string, + options map[string]any, +) (<-chan protocoltypes.StreamEvent, error) { + req, err := p.buildHTTPRequest(ctx, messages, tools, model, options, true) + if err != nil { + return nil, err + } + + resp, err := p.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("failed to send request: %w", err) + } + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() return nil, fmt.Errorf("API request failed:\n Status: %d\n Body: %s", resp.StatusCode, string(body)) } - return parseResponse(body) + ch := make(chan protocoltypes.StreamEvent, streamBufferSize) + go func() { + defer resp.Body.Close() + defer close(ch) + readSSEIntoChannel(ctx, resp.Body, ch) + }() + + return ch, nil +} + +// readSSEIntoChannel reads SSE lines from r and sends StreamEvent values on ch. +// It returns when the stream ends, an error occurs, or ctx is cancelled. +func readSSEIntoChannel(ctx context.Context, r io.Reader, ch chan<- protocoltypes.StreamEvent) { + scanner := bufio.NewScanner(r) + scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024) + + for scanner.Scan() { + // Check for context cancellation between lines. + select { + case <-ctx.Done(): + return + default: + } + + line := scanner.Text() + if !strings.HasPrefix(line, "data: ") { + continue + } + data := strings.TrimPrefix(line, "data: ") + if data == "[DONE]" { + return + } + + var chunk streamChunk + if err := json.Unmarshal([]byte(data), &chunk); err != nil { + continue // skip malformed chunks + } + + ev := protocoltypes.StreamEvent{} + + if chunk.Usage != nil { + ev.Usage = chunk.Usage + } + + if len(chunk.Choices) > 0 { + choice := chunk.Choices[0] + ev.ContentDelta = choice.Delta.Content + if choice.FinishReason != "" { + ev.FinishReason = choice.FinishReason + } + for _, tc := range choice.Delta.ToolCalls { + delta := protocoltypes.StreamToolCallDelta{ + Index: tc.Index, + ID: tc.ID, + } + if tc.Function != nil { + delta.Name = tc.Function.Name + delta.ArgumentsDelta = tc.Function.Arguments + } + ev.ToolCallDeltas = append(ev.ToolCallDeltas, delta) + } + } + + select { + case ch <- ev: + case <-ctx.Done(): + return + } + } + + if err := scanner.Err(); err != nil { + select { + case ch <- protocoltypes.StreamEvent{Err: fmt.Errorf("reading stream: %w", err)}: + case <-ctx.Done(): + } + } +} + +// AccumulateStream drains a StreamEvent channel and returns a complete LLMResponse. +func AccumulateStream(ch <-chan protocoltypes.StreamEvent) (*LLMResponse, error) { + var content strings.Builder + var toolCalls []streamToolCallAcc + var finishReason string + var usage *UsageInfo + + for ev := range ch { + if ev.Err != nil { + return nil, ev.Err + } + if ev.ContentDelta != "" { + content.WriteString(ev.ContentDelta) + } + if ev.FinishReason != "" { + finishReason = ev.FinishReason + } + if ev.Usage != nil { + usage = ev.Usage + } + for _, tc := range ev.ToolCallDeltas { + for len(toolCalls) <= tc.Index { + toolCalls = append(toolCalls, streamToolCallAcc{}) + } + if tc.ID != "" { + toolCalls[tc.Index].ID = tc.ID + } + if tc.Name != "" { + toolCalls[tc.Index].Name = tc.Name + } + toolCalls[tc.Index].Arguments.WriteString(tc.ArgumentsDelta) + } + } + + result := &LLMResponse{ + Content: content.String(), + FinishReason: finishReason, + Usage: usage, + } + + for _, tc := range toolCalls { + arguments := make(map[string]any) + argStr := tc.Arguments.String() + if argStr != "" { + if err := json.Unmarshal([]byte(argStr), &arguments); err != nil { + log.Printf("openai_compat: failed to decode streamed tool call arguments for %q: %v", tc.Name, err) + arguments["raw"] = argStr + } + } + result.ToolCalls = append(result.ToolCalls, ToolCall{ + ID: tc.ID, + Name: tc.Name, + Arguments: arguments, + }) + } + + return result, nil } func parseResponse(body []byte) (*LLMResponse, error) { @@ -240,7 +470,7 @@ func normalizeModel(model, apiBase string) string { prefix := strings.ToLower(model[:idx]) switch prefix { - case "moonshot", "nvidia", "groq", "ollama", "deepseek", "google", "openrouter", "zhipu": + case "moonshot", "nvidia", "groq", "ollama", "deepseek", "google", "openrouter", "zhipu", "minimax": return model[idx+1:] default: return model @@ -276,3 +506,129 @@ func asFloat(v any) (float64, bool) { return 0, false } } + +// --- SSE streaming support --- + +type streamChunk struct { + Choices []streamChoice `json:"choices"` + Usage *UsageInfo `json:"usage"` +} + +type streamChoice struct { + Delta streamDelta `json:"delta"` + FinishReason string `json:"finish_reason"` +} + +type streamDelta struct { + Content string `json:"content"` + ToolCalls []streamDeltaTC `json:"tool_calls"` +} + +type streamDeltaTC struct { + Index int `json:"index"` + ID string `json:"id"` + Type string `json:"type"` + Function *streamDeltaFunction `json:"function"` +} + +type streamDeltaFunction struct { + Name string `json:"name"` + Arguments string `json:"arguments"` +} + +type streamToolCallAcc struct { + ID string + Name string + Arguments strings.Builder +} + +// parseStreamResponse reads an SSE (text/event-stream) response and +// accumulates it into a single LLMResponse. +func parseStreamResponse(r io.Reader) (*LLMResponse, error) { + scanner := bufio.NewScanner(r) + // Allow up to 1 MB per SSE line to handle large argument deltas. + scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024) + + var content strings.Builder + var toolCalls []streamToolCallAcc + var finishReason string + var usage *UsageInfo + + for scanner.Scan() { + line := scanner.Text() + if !strings.HasPrefix(line, "data: ") { + continue + } + data := strings.TrimPrefix(line, "data: ") + if data == "[DONE]" { + break + } + + var chunk streamChunk + if err := json.Unmarshal([]byte(data), &chunk); err != nil { + continue // skip malformed chunks + } + + if len(chunk.Choices) == 0 { + if chunk.Usage != nil { + usage = chunk.Usage + } + continue + } + + choice := chunk.Choices[0] + if choice.Delta.Content != "" { + content.WriteString(choice.Delta.Content) + } + if choice.FinishReason != "" { + finishReason = choice.FinishReason + } + + // Accumulate streaming tool calls by index. + for _, tc := range choice.Delta.ToolCalls { + for len(toolCalls) <= tc.Index { + toolCalls = append(toolCalls, streamToolCallAcc{}) + } + if tc.ID != "" { + toolCalls[tc.Index].ID = tc.ID + } + if tc.Function != nil { + if tc.Function.Name != "" { + toolCalls[tc.Index].Name = tc.Function.Name + } + toolCalls[tc.Index].Arguments.WriteString(tc.Function.Arguments) + } + } + + if chunk.Usage != nil { + usage = chunk.Usage + } + } + if err := scanner.Err(); err != nil { + return nil, fmt.Errorf("reading stream: %w", err) + } + + result := &LLMResponse{ + Content: content.String(), + FinishReason: finishReason, + Usage: usage, + } + + for _, tc := range toolCalls { + arguments := make(map[string]any) + argStr := tc.Arguments.String() + if argStr != "" { + if err := json.Unmarshal([]byte(argStr), &arguments); err != nil { + log.Printf("openai_compat: failed to decode streamed tool call arguments for %q: %v", tc.Name, err) + arguments["raw"] = argStr + } + } + result.ToolCalls = append(result.ToolCalls, ToolCall{ + ID: tc.ID, + Name: tc.Name, + Arguments: arguments, + }) + } + + return result, nil +} diff --git a/pkg/providers/openai_compat/provider_test.go b/pkg/providers/openai_compat/provider_test.go index 42f9d42ab..7340398e9 100644 --- a/pkg/providers/openai_compat/provider_test.go +++ b/pkg/providers/openai_compat/provider_test.go @@ -1,11 +1,16 @@ package openai_compat import ( + "context" "encoding/json" + "fmt" "net/http" "net/http/httptest" "net/url" + "strings" "testing" + + "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" ) func TestProviderChat_UsesMaxCompletionTokensForGLM(t *testing.T) { @@ -281,3 +286,384 @@ func TestNormalizeModel_UsesAPIBase(t *testing.T) { t.Fatalf("normalizeModel(openrouter) = %q, want %q", got, "openrouter/auto") } } + +func TestProviderChat_StreamingTextResponse(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/text/chatcompletion_v2" { + http.Error(w, "not found", http.StatusNotFound) + return + } + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + if body["stream"] != true { + t.Error("expected stream=true in request body") + } + + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + flusher, _ := w.(http.Flusher) + + chunks := []string{ + `data: {"choices":[{"delta":{"content":"Hello"},"finish_reason":""}]}`, + `data: {"choices":[{"delta":{"content":" world"},"finish_reason":""}]}`, + `data: {"choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":2,"total_tokens":7}}`, + `data: [DONE]`, + } + for _, c := range chunks { + fmt.Fprintln(w, c) + fmt.Fprintln(w) // blank line between events + if flusher != nil { + flusher.Flush() + } + } + })) + defer server.Close() + + p := NewProviderWithOptions("key", server.URL, "", Options{ + EndpointPath: "/text/chatcompletion_v2", + Stream: true, + }) + out, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "MiniMax-M1", nil) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + if out.Content != "Hello world" { + t.Fatalf("Content = %q, want %q", out.Content, "Hello world") + } + if out.FinishReason != "stop" { + t.Fatalf("FinishReason = %q, want %q", out.FinishReason, "stop") + } + if out.Usage == nil || out.Usage.TotalTokens != 7 { + t.Fatalf("Usage.TotalTokens = %v, want 7", out.Usage) + } +} + +func TestProviderChat_StreamingToolCalls(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + flusher, _ := w.(http.Flusher) + + chunks := []string{ + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"get_weather","arguments":""}}]},"finish_reason":""}]}`, + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"city\":"}}]},"finish_reason":""}]}`, + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"SF\"}"}}]},"finish_reason":""}]}`, + `data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":10,"completion_tokens":8,"total_tokens":18}}`, + `data: [DONE]`, + } + for _, c := range chunks { + fmt.Fprintln(w, c) + fmt.Fprintln(w) + if flusher != nil { + flusher.Flush() + } + } + })) + defer server.Close() + + p := NewProviderWithOptions("key", server.URL, "", Options{Stream: true}) + out, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "weather?"}}, nil, "test", nil) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + if len(out.ToolCalls) != 1 { + t.Fatalf("len(ToolCalls) = %d, want 1", len(out.ToolCalls)) + } + tc := out.ToolCalls[0] + if tc.ID != "call_1" { + t.Fatalf("ToolCalls[0].ID = %q, want %q", tc.ID, "call_1") + } + if tc.Name != "get_weather" { + t.Fatalf("ToolCalls[0].Name = %q, want %q", tc.Name, "get_weather") + } + if tc.Arguments["city"] != "SF" { + t.Fatalf("ToolCalls[0].Arguments[city] = %v, want SF", tc.Arguments["city"]) + } +} + +func TestProviderChat_CustomEndpointPath(t *testing.T) { + var hitPath string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hitPath = r.URL.Path + resp := map[string]any{ + "choices": []map[string]any{ + {"message": map[string]any{"content": "ok"}, "finish_reason": "stop"}, + }, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) + })) + defer server.Close() + + p := NewProviderWithOptions("key", server.URL, "", Options{ + EndpointPath: "/text/chatcompletion_v2", + }) + _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "test", nil) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + if hitPath != "/text/chatcompletion_v2" { + t.Fatalf("endpoint path = %q, want %q", hitPath, "/text/chatcompletion_v2") + } +} + +func TestReadSSEIntoChannel_TextAndToolCalls(t *testing.T) { + sseData := strings.Join([]string{ + `data: {"choices":[{"delta":{"content":"Hello"},"finish_reason":""}]}`, + ``, + `data: {"choices":[{"delta":{"content":" world"},"finish_reason":""}]}`, + ``, + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"greet","arguments":"{\"n"}}]},"finish_reason":""}]}`, + ``, + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"ame\":\"Bob\"}"}}]},"finish_reason":""}]}`, + ``, + `data: {"choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":4,"total_tokens":7}}`, + ``, + `data: [DONE]`, + ``, + }, "\n") + + ch := make(chan protocoltypes.StreamEvent, 32) + go func() { + defer close(ch) + readSSEIntoChannel(context.Background(), strings.NewReader(sseData), ch) + }() + + var events []protocoltypes.StreamEvent + for ev := range ch { + events = append(events, ev) + } + + if len(events) < 3 { + t.Fatalf("got %d events, want at least 3", len(events)) + } + + // Check content deltas + if events[0].ContentDelta != "Hello" { + t.Errorf("events[0].ContentDelta = %q, want %q", events[0].ContentDelta, "Hello") + } + if events[1].ContentDelta != " world" { + t.Errorf("events[1].ContentDelta = %q, want %q", events[1].ContentDelta, " world") + } + + // Check tool call deltas + if len(events[2].ToolCallDeltas) != 1 || events[2].ToolCallDeltas[0].ID != "call_1" { + t.Errorf("events[2] should contain tool call with ID=call_1") + } + if events[2].ToolCallDeltas[0].Name != "greet" { + t.Errorf("events[2].ToolCallDeltas[0].Name = %q, want %q", events[2].ToolCallDeltas[0].Name, "greet") + } + + // Check finish event + lastEv := events[len(events)-1] + if lastEv.FinishReason != "stop" { + t.Errorf("last event FinishReason = %q, want %q", lastEv.FinishReason, "stop") + } + if lastEv.Usage == nil || lastEv.Usage.TotalTokens != 7 { + t.Errorf("last event Usage.TotalTokens = %v, want 7", lastEv.Usage) + } +} + +func TestReadSSEIntoChannel_ContextCancel(t *testing.T) { + // Simulate a slow SSE stream that gets cancelled. + ctx, cancel := context.WithCancel(context.Background()) + + // Create a reader that blocks after sending one chunk. + sseData := `data: {"choices":[{"delta":{"content":"first"},"finish_reason":""}]}` + "\n\n" + + ch := make(chan protocoltypes.StreamEvent, 32) + go func() { + defer close(ch) + readSSEIntoChannel(ctx, strings.NewReader(sseData), ch) + }() + + // Read the first event. + ev := <-ch + if ev.ContentDelta != "first" { + t.Fatalf("ContentDelta = %q, want %q", ev.ContentDelta, "first") + } + + // Cancel the context; the channel should close. + cancel() + _, ok := <-ch + if ok { + t.Fatal("expected channel to be closed after context cancel") + } +} + +func TestAccumulateStream_FullResponse(t *testing.T) { + ch := make(chan protocoltypes.StreamEvent, 8) + + go func() { + ch <- protocoltypes.StreamEvent{ContentDelta: "Hello"} + ch <- protocoltypes.StreamEvent{ContentDelta: " world"} + ch <- protocoltypes.StreamEvent{ + ToolCallDeltas: []protocoltypes.StreamToolCallDelta{ + {Index: 0, ID: "call_1", Name: "test_tool", ArgumentsDelta: `{"key"`}, + }, + } + ch <- protocoltypes.StreamEvent{ + ToolCallDeltas: []protocoltypes.StreamToolCallDelta{ + {Index: 0, ArgumentsDelta: `:"value"}`}, + }, + } + ch <- protocoltypes.StreamEvent{ + FinishReason: "stop", + Usage: &UsageInfo{PromptTokens: 5, CompletionTokens: 3, TotalTokens: 8}, + } + close(ch) + }() + + resp, err := AccumulateStream(ch) + if err != nil { + t.Fatalf("AccumulateStream() error = %v", err) + } + + if resp.Content != "Hello world" { + t.Errorf("Content = %q, want %q", resp.Content, "Hello world") + } + if resp.FinishReason != "stop" { + t.Errorf("FinishReason = %q, want %q", resp.FinishReason, "stop") + } + if resp.Usage == nil || resp.Usage.TotalTokens != 8 { + t.Errorf("Usage.TotalTokens = %v, want 8", resp.Usage) + } + if len(resp.ToolCalls) != 1 { + t.Fatalf("len(ToolCalls) = %d, want 1", len(resp.ToolCalls)) + } + if resp.ToolCalls[0].Name != "test_tool" { + t.Errorf("ToolCalls[0].Name = %q, want %q", resp.ToolCalls[0].Name, "test_tool") + } + if resp.ToolCalls[0].Arguments["key"] != "value" { + t.Errorf("ToolCalls[0].Arguments[key] = %v, want %q", resp.ToolCalls[0].Arguments["key"], "value") + } +} + +func TestAccumulateStream_Error(t *testing.T) { + ch := make(chan protocoltypes.StreamEvent, 4) + + go func() { + ch <- protocoltypes.StreamEvent{ContentDelta: "partial"} + ch <- protocoltypes.StreamEvent{Err: fmt.Errorf("connection reset")} + close(ch) + }() + + _, err := AccumulateStream(ch) + if err == nil { + t.Fatal("expected error, got nil") + } + if !strings.Contains(err.Error(), "connection reset") { + t.Fatalf("error = %q, want to contain %q", err.Error(), "connection reset") + } +} + +func TestChatStream_EndToEnd(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + flusher, _ := w.(http.Flusher) + + chunks := []string{ + `data: {"choices":[{"delta":{"content":"stream"},"finish_reason":""}]}`, + `data: {"choices":[{"delta":{"content":"ed"},"finish_reason":""}]}`, + `data: {"choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":2,"completion_tokens":1,"total_tokens":3}}`, + `data: [DONE]`, + } + for _, c := range chunks { + fmt.Fprintln(w, c) + fmt.Fprintln(w) + if flusher != nil { + flusher.Flush() + } + } + })) + defer server.Close() + + p := NewProviderWithOptions("key", server.URL, "", Options{Stream: true}) + + ch, err := p.ChatStream(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "test", nil) + if err != nil { + t.Fatalf("ChatStream() error = %v", err) + } + + resp, err := AccumulateStream(ch) + if err != nil { + t.Fatalf("AccumulateStream() error = %v", err) + } + + if resp.Content != "streamed" { + t.Errorf("Content = %q, want %q", resp.Content, "streamed") + } + if resp.FinishReason != "stop" { + t.Errorf("FinishReason = %q, want %q", resp.FinishReason, "stop") + } + if resp.Usage == nil || resp.Usage.TotalTokens != 3 { + t.Errorf("Usage.TotalTokens = %v, want 3", resp.Usage) + } +} + +func TestChatStream_EarlyCancel(t *testing.T) { + serverDone := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer close(serverDone) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + flusher, _ := w.(http.Flusher) + + // Send many chunks; expect the client to cancel early. + for i := 0; i < 1000; i++ { + select { + case <-r.Context().Done(): + return + default: + } + fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":\"x\"},\"finish_reason\":\"\"}]}\n\n") + if flusher != nil { + flusher.Flush() + } + } + })) + defer server.Close() + + p := NewProviderWithOptions("key", server.URL, "", Options{Stream: true}) + + ctx, cancel := context.WithCancel(context.Background()) + ch, err := p.ChatStream(ctx, []Message{{Role: "user", Content: "hi"}}, nil, "test", nil) + if err != nil { + t.Fatalf("ChatStream() error = %v", err) + } + + // Read a few events, then cancel. + count := 0 + for ev := range ch { + if ev.Err != nil { + break + } + count++ + if count >= 5 { + cancel() + } + } + + if count < 5 { + t.Errorf("expected at least 5 events before cancel, got %d", count) + } + + // Server should have received the cancellation. + <-serverDone +} + +func TestCanStream(t *testing.T) { + p1 := NewProvider("key", "https://example.com", "") + if p1.CanStream() { + t.Error("CanStream() = true for non-stream provider") + } + + p2 := NewProviderWithOptions("key", "https://example.com", "", Options{Stream: true}) + if !p2.CanStream() { + t.Error("CanStream() = false for stream provider") + } +} diff --git a/pkg/providers/protocoltypes/types.go b/pkg/providers/protocoltypes/types.go index 3a089ca47..65d7a2c05 100644 --- a/pkg/providers/protocoltypes/types.go +++ b/pkg/providers/protocoltypes/types.go @@ -54,3 +54,20 @@ type ToolFunctionDefinition struct { Description string `json:"description"` Parameters map[string]any `json:"parameters"` } + +// StreamEvent represents a single chunk from an SSE streaming response. +type StreamEvent struct { + ContentDelta string + ToolCallDeltas []StreamToolCallDelta + FinishReason string // set only on the final event + Usage *UsageInfo // set only on the final event + Err error // non-nil when the stream encountered an error +} + +// StreamToolCallDelta carries an incremental piece of a streaming tool call. +type StreamToolCallDelta struct { + Index int + ID string // set on the first chunk for this tool call + Name string // set on the first chunk for this tool call + ArgumentsDelta string // JSON fragment (incremental) +} diff --git a/pkg/providers/types.go b/pkg/providers/types.go index f711e7803..bdce83765 100644 --- a/pkg/providers/types.go +++ b/pkg/providers/types.go @@ -17,6 +17,8 @@ type ( ToolFunctionDefinition = protocoltypes.ToolFunctionDefinition ExtraContent = protocoltypes.ExtraContent GoogleExtra = protocoltypes.GoogleExtra + StreamEvent = protocoltypes.StreamEvent + StreamToolCallDelta = protocoltypes.StreamToolCallDelta ) type LLMProvider interface { @@ -67,6 +69,22 @@ func (e *FailoverError) IsRetriable() bool { return e.Reason != FailoverFormat } +// StreamingProvider extends LLMProvider with SSE channel-based streaming. +// Use a type assertion to check if a provider supports streaming: +// +// if sp, ok := provider.(StreamingProvider); ok && sp.CanStream() { ... } +type StreamingProvider interface { + LLMProvider + CanStream() bool + ChatStream( + ctx context.Context, + messages []Message, + tools []ToolDefinition, + model string, + options map[string]any, + ) (<-chan StreamEvent, error) +} + // ModelConfig holds primary model and fallback list. type ModelConfig struct { Primary string diff --git a/pkg/tools/shell.go b/pkg/tools/shell.go index ba2290f37..e7c04fed4 100644 --- a/pkg/tools/shell.go +++ b/pkg/tools/shell.go @@ -291,10 +291,16 @@ func (t *ExecTool) guardCommand(command, cwd string) string { } // Token-based absolute path detection. - // Uses strings.Fields instead of regex to avoid false positives - // from slashes in relative paths (e.g., "tests/cold/file.py"). + // Uses strings.Fields so relative paths (e.g., "tests/cold/file.py") + // are not falsely flagged. // Flags like -I/usr/local/include are naturally skipped because // filepath.IsAbs returns false for tokens starting with "-". + // + // Agent CLI tools (claude, codex, gemini) accept slash commands + // (e.g., "/review") that look like absolute paths but are not. + // For these tools we check whether the token is an existing path + // before blocking. + agentCLI := isAgentCLICommand(cmd) for _, token := range strings.Fields(cmd) { token = strings.Trim(token, "\"'") @@ -313,6 +319,13 @@ func (t *ExecTool) guardCommand(command, cwd string) string { if isExecutable(p) { continue } + // Agent CLI slash commands: skip non-existent paths + // (e.g., "/review" is a command, not a file). + if agentCLI { + if _, statErr := os.Stat(p); os.IsNotExist(statErr) { + continue + } + } return "Command blocked by safety guard (path outside working dir)" } } @@ -321,6 +334,25 @@ func (t *ExecTool) guardCommand(command, cwd string) string { return "" } +// agentCLINames lists agent CLI tools that use slash commands +// (e.g., "/review", "/help") which look like absolute paths. +var agentCLINames = []string{"claude", "codex", "gemini"} + +// isAgentCLICommand returns true if the command invokes an agent CLI tool. +func isAgentCLICommand(cmd string) bool { + fields := strings.Fields(cmd) + if len(fields) == 0 { + return false + } + base := filepath.Base(fields[0]) + for _, name := range agentCLINames { + if base == name { + return true + } + } + return false +} + // isExecutable checks if a path points to an executable file. // On Unix, checks the execute permission bits. // On Windows, checks for known executable extensions. diff --git a/pkg/tools/shell_test.go b/pkg/tools/shell_test.go index 9a023ae68..d760af5b4 100644 --- a/pkg/tools/shell_test.go +++ b/pkg/tools/shell_test.go @@ -486,3 +486,31 @@ func TestGuardCommand_CdWithAbsoluteWorkspacePath(t *testing.T) { t.Errorf("cd to workspace subdir should be allowed: %q → %s", cmd, result) } } + +func TestGuardCommand_AgentCLISlashCommand(t *testing.T) { + workspace := t.TempDir() + tool := NewExecTool(workspace, true) + + // Agent CLI slash commands (e.g., "/review") are not file paths. + // They should be allowed because they don't exist on disk. + cmds := []string{ + `codex exec --yolo "/review skip-git-repo-check"`, + `claude "/review"`, + `gemini "/help"`, + } + for _, cmd := range cmds { + result := tool.guardCommand(cmd, workspace) + if result != "" { + t.Errorf("Agent CLI slash command should not be blocked: %q → %s", cmd, result) + } + } + + // Non-agent commands with absolute paths should still be blocked. + if runtime.GOOS != "windows" { + blocked := `cat /etc/hosts` + result := tool.guardCommand(blocked, workspace) + if result == "" { + t.Errorf("Non-agent command with absolute path should be blocked: %q", blocked) + } + } +} diff --git a/pkg/utils/string.go b/pkg/utils/string.go index 7a6aa37cc..df1369dae 100644 --- a/pkg/utils/string.go +++ b/pkg/utils/string.go @@ -1,5 +1,95 @@ package utils +import ( + "regexp" + "strings" +) + +// Repetition detection constants. +const ( + repetitionSampleSize = 2000 // runes to sample from the tail + repetitionNgramSize = 10 // sliding window length + repetitionUniqueThreshold = 0.1 // unique ratio below this → repetition +) + +var ( + thinkBlockClosedRe = regexp.MustCompile(`(?is).*?`) + thinkBlockOpenRe = regexp.MustCompile(`(?is).*$`) +) + +// StripThinkBlocks removes blocks (including unclosed ones) +// from s and returns the trimmed result. +func StripThinkBlocks(s string) string { + s = thinkBlockClosedRe.ReplaceAllString(s, "") + s = thinkBlockOpenRe.ReplaceAllString(s, "") + return strings.TrimSpace(s) +} + +// TailPad returns a fixed-height block of n visual lines built from the +// tail of s. Long lines are wrapped at wrapWidth runes so the result +// never exceeds the chat bubble width. If fewer than n visual lines +// exist, Braille-blank lines (\u2800) are prepended as padding. +func TailPad(s string, n, wrapWidth int) string { + // Wrap each raw line into visual lines respecting wrapWidth. + var visual []string + for _, raw := range strings.Split(s, "\n") { + visual = append(visual, wrapLine(raw, wrapWidth)...) + } + if len(visual) > n { + visual = visual[len(visual)-n:] + } + for len(visual) < n { + visual = append([]string{"\u2800"}, visual...) + } + return strings.Join(visual, "\n") +} + +// wrapLine splits a single line into segments of at most width runes. +// An empty line produces one empty string (preserving blank lines). +func wrapLine(line string, width int) []string { + runes := []rune(line) + if len(runes) <= width { + return []string{line} + } + var segs []string + for len(runes) > 0 { + end := width + if end > len(runes) { + end = len(runes) + } + segs = append(segs, string(runes[:end])) + runes = runes[end:] + } + return segs +} + +// DetectRepetitionLoop checks if text contains degenerate repetition +// by computing the unique N-gram ratio on the last repetitionSampleSize runes. +// Returns true if the ratio of unique N-grams to total N-grams +// falls below repetitionUniqueThreshold (i.e., 90%+ are duplicates). +func DetectRepetitionLoop(text string) bool { + runes := []rune(text) + + // Sample the tail + if len(runes) > repetitionSampleSize { + runes = runes[len(runes)-repetitionSampleSize:] + } + + total := len(runes) - repetitionNgramSize + 1 + if total <= 0 { + return false + } + + unique := make(map[string]struct{}, total/repetitionNgramSize) + for i := 0; i < total; i++ { + ng := string(runes[i : i+repetitionNgramSize]) + unique[ng] = struct{}{} + } + + ratio := float64(len(unique)) / float64(total) + return ratio < repetitionUniqueThreshold +} + // Truncate returns a truncated version of s with at most maxLen runes. // Handles multi-byte Unicode characters properly. // If the string is truncated, "..." is appended to indicate truncation. diff --git a/pkg/utils/string_test.go b/pkg/utils/string_test.go new file mode 100644 index 000000000..bde040384 --- /dev/null +++ b/pkg/utils/string_test.go @@ -0,0 +1,197 @@ +package utils + +import ( + "strings" + "testing" +) + +// --- StripThinkBlocks --- + +func TestStripThinkBlocks_ClosedBlock(t *testing.T) { + in := "\nsecret reasoning\n\n\nVisible content" + got := StripThinkBlocks(in) + if got != "Visible content" { + t.Fatalf("StripThinkBlocks() = %q, want %q", got, "Visible content") + } +} + +func TestStripThinkBlocks_UnclosedBlock(t *testing.T) { + in := "reasoning that never ends\nmore reasoning" + got := StripThinkBlocks(in) + if got != "" { + t.Fatalf("StripThinkBlocks() = %q, want empty", got) + } +} + +func TestStripThinkBlocks_MultipleBlocks(t *testing.T) { + in := "firstmiddlesecondend" + got := StripThinkBlocks(in) + if got != "middleend" { + t.Fatalf("StripThinkBlocks() = %q, want %q", got, "middleend") + } +} + +func TestStripThinkBlocks_NoBlocks(t *testing.T) { + in := "plain text without think blocks" + got := StripThinkBlocks(in) + if got != in { + t.Fatalf("StripThinkBlocks() = %q, want %q", got, in) + } +} + +func TestStripThinkBlocks_CaseInsensitive(t *testing.T) { + in := "upper casevisible" + got := StripThinkBlocks(in) + if got != "visible" { + t.Fatalf("StripThinkBlocks() = %q, want %q", got, "visible") + } +} + +func TestStripThinkBlocks_ClosedThenUnclosed(t *testing.T) { + in := "closedmiddleunclosed tail" + got := StripThinkBlocks(in) + if got != "middle" { + t.Fatalf("StripThinkBlocks() = %q, want %q", got, "middle") + } +} + +// --- DetectRepetitionLoop --- + +func TestDetectRepetitionLoop_HighRepetition(t *testing.T) { + // Repeat a short phrase many times → should be detected + phrase := "結構本格的なコード" + repeated := strings.Repeat(phrase, 300) + if !DetectRepetitionLoop(repeated) { + t.Fatal("DetectRepetitionLoop should return true for highly repetitive text") + } +} + +func TestDetectRepetitionLoop_NormalText(t *testing.T) { + // Normal varied text should not trigger + normal := "The quick brown fox jumps over the lazy dog. " + + "Pack my box with five dozen liquor jugs. " + + "How vexingly quick daft zebras jump. " + + "Sphinx of black quartz, judge my vow. " + + "Two driven jocks help fax my big quiz. " + + "The five boxing wizards jump quickly. " + + "Jackdaws love my big sphinx of quartz. " + + "Grumpy wizards make a toxic brew for the jovial queen." + // Extend to be long enough + long := strings.Repeat(normal+" ", 10) + if DetectRepetitionLoop(long) { + t.Fatal("DetectRepetitionLoop should return false for normal text") + } +} + +func TestDetectRepetitionLoop_ShortText(t *testing.T) { + // Text shorter than N-gram size should never trigger + if DetectRepetitionLoop("short") { + t.Fatal("DetectRepetitionLoop should return false for short text") + } +} + +func TestDetectRepetitionLoop_EmptyString(t *testing.T) { + if DetectRepetitionLoop("") { + t.Fatal("DetectRepetitionLoop should return false for empty string") + } +} + +func TestDetectRepetitionLoop_SingleCharRepeat(t *testing.T) { + // "aaaa..." repeated → only 1 unique N-gram → detected + repeated := strings.Repeat("あ", 2500) + if !DetectRepetitionLoop(repeated) { + t.Fatal("DetectRepetitionLoop should return true for single-char repetition") + } +} + +func TestDetectRepetitionLoop_BelowSampleSize(t *testing.T) { + // Repetitive but under sample size still detected + phrase := "abcdefghij" + repeated := strings.Repeat(phrase, 50) // 500 chars + if !DetectRepetitionLoop(repeated) { + t.Fatal("DetectRepetitionLoop should return true for repetitive text below sample size") + } +} + +// --- TailPad --- + +func TestTailPad_FewerThanN(t *testing.T) { + got := TailPad("a\nb", 5, 80) + lines := strings.Split(got, "\n") + if len(lines) != 5 { + t.Fatalf("TailPad line count = %d, want 5", len(lines)) + } + for i := 0; i < 3; i++ { + if lines[i] != "\u2800" { + t.Errorf("TailPad line %d = %q, want padding", i, lines[i]) + } + } + if lines[3] != "a" || lines[4] != "b" { + t.Errorf("TailPad content = %q %q, want a b", lines[3], lines[4]) + } +} + +func TestTailPad_ExactlyN(t *testing.T) { + in := "a\nb\nc" + got := TailPad(in, 3, 80) + if got != in { + t.Fatalf("TailPad exact = %q, want %q", got, in) + } +} + +func TestTailPad_MoreThanN(t *testing.T) { + got := TailPad("a\nb\nc\nd\ne", 3, 80) + if got != "c\nd\ne" { + t.Fatalf("TailPad tail = %q, want %q", got, "c\nd\ne") + } +} + +func TestTailPad_Empty(t *testing.T) { + got := TailPad("", 4, 80) + lines := strings.Split(got, "\n") + if len(lines) != 4 { + t.Fatalf("TailPad empty line count = %d, want 4", len(lines)) + } + for i, l := range lines { + if i == len(lines)-1 { + if l != "" { + t.Errorf("TailPad empty last line = %q, want empty", l) + } + } else if l != "\u2800" { + t.Errorf("TailPad empty line %d = %q, want padding", i, l) + } + } +} + +func TestTailPad_LongLineWraps(t *testing.T) { + // One 10-char line wraps into 2 visual lines at width 5. + got := TailPad("abcdefghij", 4, 5) + lines := strings.Split(got, "\n") + if len(lines) != 4 { + t.Fatalf("TailPad wrap line count = %d, want 4", len(lines)) + } + // 2 padding + "abcde" + "fghij" + if lines[2] != "abcde" || lines[3] != "fghij" { + t.Errorf("TailPad wrap content = %v", lines) + } +} + +func TestTailPad_WrapPushesOldLines(t *testing.T) { + // "short" (1 visual) + "abcdefghij" (2 visual at width 5) = 3 visual. + // With n=2, only tail 2 visual lines remain. + got := TailPad("short\nabcdefghij", 2, 5) + if got != "abcde\nfghij" { + t.Fatalf("TailPad wrap push = %q, want %q", got, "abcde\nfghij") + } +} + +// --- Truncate --- + +func TestTruncate(t *testing.T) { + if got := Truncate("hello", 10); got != "hello" { + t.Errorf("Truncate short = %q", got) + } + if got := Truncate("hello world!", 8); got != "hello..." { + t.Errorf("Truncate long = %q", got) + } +}