Merge remote-tracking branch 'origin/main' into feature/miniapp-git-tab
This commit is contained in:
dj-oyu
commit
c90bb5d3c8
18 changed files with 1780 additions and 37 deletions
|
|
@ -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",
|
||||
|
|
|
|||
112
docs/plan-interview-improvements.md
Normal file
112
docs/plan-interview-improvements.md
Normal file
|
|
@ -0,0 +1,112 @@
|
|||
# /plan interview 改善検討
|
||||
|
||||
## 現象
|
||||
|
||||
`/plan <task>` でinterviewモードに入った後、AIが:
|
||||
|
||||
1. MEMORY.mdに何も書かずに会話だけ続ける
|
||||
2. interviewを無視して実装を始めようとする（exec, ファイル書き込み）
|
||||
3. 結局フェーズ/ステップ/コマンドが書かれないまま executing に遷移する
|
||||
|
||||
## 原因分析
|
||||
|
||||
### Layer 1: 技術的障壁（修正済み）
|
||||
|
||||
| 問題 | 原因 | 修正 |
|
||||
|---|---|---|
|
||||
| XML tool callがパースされない | MiniMaxの開閉タグ不一致（`<minimax:toolcall>` vs `</minimax:tool_call>`） | 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 <task> → 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が補助するが確実ではない
|
||||
|
|
@ -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",
|
||||
|
|
@ -1240,6 +1249,7 @@ 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"
|
||||
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 <think> 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 {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -51,8 +51,6 @@ const telegramMaxMessageChars = 3900
|
|||
const markdownTableMaxWidth = 42
|
||||
const markdownTableMinColWidth = 6
|
||||
|
||||
var thinkBlockPattern = regexp.MustCompile(`(?is)<think>.*?</think>`)
|
||||
|
||||
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)"
|
||||
}
|
||||
|
|
|
|||
|
|
@ -5,8 +5,8 @@ import (
|
|||
"testing"
|
||||
)
|
||||
|
||||
func TestSanitizeTelegramOutgoingContent_RemovesThinkBlock(t *testing.T) {
|
||||
in := "<think>\nsecret reasoning\n</think>\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 := "<think>only reasoning</think>"
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)<think>.*?</think>`)
|
||||
thinkBlockOpenRe = regexp.MustCompile(`(?is)<think>.*$`)
|
||||
)
|
||||
|
||||
// StripThinkBlocks removes <think>…</think> 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.
|
||||
|
|
|
|||
197
pkg/utils/string_test.go
Normal file
197
pkg/utils/string_test.go
Normal file
|
|
@ -0,0 +1,197 @@
|
|||
package utils
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// --- StripThinkBlocks ---
|
||||
|
||||
func TestStripThinkBlocks_ClosedBlock(t *testing.T) {
|
||||
in := "<think>\nsecret reasoning\n</think>\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 := "<think>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 := "<think>first</think>middle<think>second</think>end"
|
||||
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 := "<THINK>upper case</THINK>visible"
|
||||
got := StripThinkBlocks(in)
|
||||
if got != "visible" {
|
||||
t.Fatalf("StripThinkBlocks() = %q, want %q", got, "visible")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStripThinkBlocks_ClosedThenUnclosed(t *testing.T) {
|
||||
in := "<think>closed</think>middle<think>unclosed 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)
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue