diff --git a/Makefile b/Makefile index 152f724..176a95d 100644 --- a/Makefile +++ b/Makefile @@ -1,6 +1,6 @@ # dsfree2api — build / test helpers GO ?= go -VERSION ?= 0.3.0 +VERSION ?= 0.4.0 LDFLAGS := -s -w -X github.com/nyoungo/dsfree2api/internal/admin.Version=$(VERSION) BIN := bin/dsfree2api IMAGE ?= dsfree2api:latest diff --git a/README.md b/README.md index eb32bca..ccc96bc 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,6 @@ Go 实现，编译为**单个静态二进制**，无运行时依赖；内置 **W - [API 用法](#api-用法) - [Web 管理台](#web-管理台) - [环境变量](#环境变量) -- [与 Python 版的差异](#与-python-版的差异) - [架构](#架构) - [开发](#开发) - [免责声明](#免责声明) @@ -32,7 +31,7 @@ Go 实现，编译为**单个静态二进制**，无运行时依赖；内置 **W ## 功能特性 - **多格式对话接口**：`POST /v1/chat/completions`（OpenAI，流式 / 非流式）、`POST /v1/responses`（OpenAI Responses）、`POST /v1/messages`（Anthropic Messages）、`GET /v1/models`、`GET /health` -- **Tool Calling**：三个对话端点均支持 `tools` / `tool_choice`，模型以 JSON 输出 tool_calls，自动回填到响应（`tool_calls` / `function_call` / `tool_use`） +- **Tool Calling**：三个对话端点均支持 `tools` / `tool_choice`，模型以 JSON 输出 tool_calls，自动回填到响应（`tool_calls` / `function_call` / `tool_use`）；回复被上游输出上限截断时自动多轮续写拼接（`continue_rounds`） - **三站点聚合**：单站点配额耗尽自动切换到其他站点的同款模型 - **代理线路**：主线路 + 多条备用线路，支持 `http://` `https://` `socks5://`；slow-start 检测 + 失败降级 + 流中失败不重放 - **Turnstile 自动突破**：三种 Token 获取方式 — 求解服务 API（`provider="api"`）、本地浏览器自动求解（`provider="browser"`，CDP 直连 Chrome/Edge，无需 Playwright/Docker）、手动导入 Cookie（`provider="manual"`）；结果按「代理线路 × 站点」缓存，TTL 过期自动重解 @@ -139,7 +138,7 @@ dsfree2api -config config.toml -check | `[admin]` | 管理台 `enabled` / `host` / `port` / `password`（留空启动时生成随机密码） | | `[limits]` | `max_concurrent_per_site` 每站点并发、`rate_per_minute` 每 Key 每分钟限流（0 = 不限） | | `[proxy]` | `url` 主线路、`fallback_urls` 备用线路、`slow_start_seconds` 首事件超时 | -| `[upstream]` | 超时、配置缓存 TTL、会话自动刷新、`cross_site_failover` 跨站切换 | +| `[upstream]` | 超时、配置缓存 TTL、会话自动刷新、`cross_site_failover` 跨站切换、`continue_rounds` 工具调用截断续写轮数 | | `[turnstile]` | **默认 `enabled = false` 且不内置任何求解服务**；`provider` 三选一：`api`（填 `api_url` / `api_key`）、`browser`（填 `browser_path` 指向本机 Chrome/Edge）、`manual`（只用控制台导入的 Cookie），另有 Cookie TTL、重试次数 | | `[sites.*]` | 三个站点的 `base_url` / `ajax_url` / `sitekey` / `language` | | `[models.*]` | 模型到 `site` + `upstream_id` + `bot_id` / `post_id` 的映射 | @@ -296,30 +295,6 @@ curl http://127.0.0.1:8000/v1/messages \ --- -## 与 Python 版的差异 - -功能对齐 `deepseek-fr-2api-main`（Python/FastAPI 版），并做了以下修复与增强： - -**修复** - -- 流式 `tool_calls` 仅在请求带 `tools` 时缓冲后再下发，避免内容重复 -- `/v1/responses` 流式错误发 `response.failed`，不再假报 `completed` -- 会话标记识别只匹配明确 marker，不再误判裸 `"nonce"` 字段 -- `temperature` / `max_tokens` / `stop` / `response_format` 以 prompt 尾部约束下发 -- 空 role 不再导致 prompt 构建 panic - -**增强** - -- 站点配额耗尽跨站切换（`cross_site_failover`） -- 每站点并发闸门 + 每 Key 每分钟限流 -- 启动期配置校验（`-check`），按 provider 检查 `api_url` / `api_key` 或 `browser_path`，缺失直接报错而非运行时失败 -- 统计持久化（`data/metrics.json`）、实时日志 SSE、对话调试页 -- `POST /v1/messages` Anthropic Messages 格式（流式 / 非流式 / `tool_use`），`/v1/responses` 支持 `tools` → `function_call`，多轮 tool 结果回填 -- `Chat` / `Responses` / `Messages` / `Models` / 代理与 Turnstile 的单元测试 -- 单二进制部署，无 Python / uv / libcurl 依赖 - ---- - ## 架构 ``` diff --git a/config.example.toml b/config.example.toml index 5b27c5d..c9cbd3b 100644 --- a/config.example.toml +++ b/config.example.toml @@ -43,6 +43,9 @@ retry_backoff_seconds = 1.0 user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/149.0.0.0 Safari/537.36" # 本站点配额耗尽时自动切换到其他站点的同款模型。 cross_site_failover = true +# 站点对单次回复有约 1k token 的硬上限，大参数的工具调用 JSON 会被截断。 +# 自动向上游追加「续写」轮次把 JSON 拼完整；0 = 关闭。 +continue_rounds = 20 [turnstile] # 默认关闭，且不内置任何求解服务。Token 获取方式三选一（provider）： diff --git a/internal/admin/server.go b/internal/admin/server.go index 269eaa2..847e917 100644 --- a/internal/admin/server.go +++ b/internal/admin/server.go @@ -20,7 +20,7 @@ import ( ) // Version is reported by the console and /api/overview. -var Version = "0.3.0" +var Version = "0.4.0" //go:embed web/* var webFS embed.FS diff --git a/internal/api/chat.go b/internal/api/chat.go index 1f3a933..0d2c038 100644 --- a/internal/api/chat.go +++ b/internal/api/chat.go @@ -64,11 +64,12 @@ func (s *Server) handleChat(w http.ResponseWriter, r *http.Request) { s.met.IncInFlight(1) defer s.met.IncInFlight(-1) + cont := continuationFor(req.Tools, req.Messages, s.continueRounds()) if req.Stream { - s.streamChat(ctx, w, req, prompt, site) + s.streamChat(ctx, w, req, prompt, site, cont) return } - s.completeChat(ctx, w, req, prompt, site) + s.completeChat(ctx, w, req, prompt, site, cont) } // chatOutcome is the result of replaying one upstream conversation. @@ -84,20 +85,21 @@ type chatOutcome struct { } // collect runs a full (non-streaming) conversation against the upstream site -// and records metrics for it. -func (s *Server) collect(ctx context.Context, modelID, prompt, site string) chatOutcome { +// and records metrics for it. When cont is non-nil the reply is automatically +// continued across extra upstream turns until a truncated tool-call JSON +// closes (see openai.RunContinued). +func (s *Server) collect(ctx context.Context, modelID, prompt, site string, cont openai.ContinuationFunc) chatOutcome { start := time.Now() info := &upstream.ServeInfo{} var text strings.Builder var firstDelta time.Time - err := s.up.Chat(ctx, modelID, prompt, info, func(ev upstream.Event) error { - if ev.Kind == upstream.KindDelta && ev.Value != "" { - if firstDelta.IsZero() { - firstDelta = time.Now() - } - text.WriteString(ev.Value) + err := s.chatWithContinue(ctx, modelID, prompt, cont, info, func() { + if firstDelta.IsZero() { + firstDelta = time.Now() } + }, func(seg string) error { + text.WriteString(seg) return nil }) @@ -132,8 +134,8 @@ func (s *Server) collect(ctx context.Context, modelID, prompt, site string) chat } } -func (s *Server) completeChat(ctx context.Context, w http.ResponseWriter, req openai.ChatCompletionRequest, prompt, site string) { - out := s.collect(ctx, req.Model, prompt, site) +func (s *Server) completeChat(ctx context.Context, w http.ResponseWriter, req openai.ChatCompletionRequest, prompt, site string, cont openai.ContinuationFunc) { + out := s.collect(ctx, req.Model, prompt, site, cont) if out.err != nil { writeUpstreamError(w, out.err) return @@ -175,7 +177,7 @@ func finishReason(calls []openai.ToolCall) string { return "stop" } -func (s *Server) streamChat(ctx context.Context, w http.ResponseWriter, req openai.ChatCompletionRequest, prompt, site string) { +func (s *Server) streamChat(ctx context.Context, w http.ResponseWriter, req openai.ChatCompletionRequest, prompt, site string, cont openai.ContinuationFunc) { flusher, _ := w.(http.Flusher) w.Header().Set("Content-Type", "text/event-stream; charset=utf-8") w.Header().Set("Cache-Control", "no-cache") @@ -221,18 +223,16 @@ func (s *Server) streamChat(ctx context.Context, w http.ResponseWriter, req open s.met.IncInFlight(1) defer s.met.IncInFlight(-1) - err := s.up.Chat(ctx, req.Model, prompt, info, func(ev upstream.Event) error { - if ev.Kind != upstream.KindDelta || ev.Value == "" { - return nil - } + err := s.chatWithContinue(ctx, req.Model, prompt, cont, info, func() { if firstDelta.IsZero() { firstDelta = time.Now() } - text.WriteString(ev.Value) + }, func(seg string) error { + text.WriteString(seg) if buffered { return nil } - return write(chunk(openai.ChoiceDelta{Content: ev.Value}, nil)) + return write(chunk(openai.ChoiceDelta{Content: seg}, nil)) }) duration := time.Since(start) diff --git a/internal/api/continue.go b/internal/api/continue.go new file mode 100644 index 0000000..d589c98 --- /dev/null +++ b/internal/api/continue.go @@ -0,0 +1,62 @@ +package api + +import ( + "context" + + "github.com/nyoungo/dsfree2api/internal/openai" + "github.com/nyoungo/dsfree2api/internal/upstream" +) + +// continuationFor builds the prompt function that asks the model to resume a +// tool-call reply the site's per-reply output cap cut off. It returns nil +// when no tools were offered, which disables continuation entirely. +func continuationFor(tools []openai.ToolDef, messages []openai.ChatMessage, rounds int) openai.ContinuationFunc { + if len(tools) == 0 || len(messages) == 0 { + return nil + } + base := append([]openai.ChatMessage{}, messages...) + return func(partial string, round int) string { + history := append(append([]openai.ChatMessage{}, base...), + openai.ChatMessage{Role: "assistant", Content: openai.MessageContent{Text: openai.ContinuationContext(partial, 400, 2000)}}, + openai.ChatMessage{Role: "user", Content: openai.MessageContent{Text: openai.ContinueInstruction(partial, round, rounds)}}, + ) + return openai.BuildPrompt(history, openai.PromptOptions{}) + } +} + +// continueRounds reads the configured number of continuation turns. +func (s *Server) continueRounds() int { + s.cfg.RLock() + defer s.cfg.RUnlock() + return s.cfg.Upstream.ContinueRounds +} + +// chatWithContinue runs one upstream conversation and transparently stitches +// truncated tool-call JSON replies across extra continuation turns. emit +// receives the reply text in order; firstDelta is stamped on the first delta +// of any turn. +func (s *Server) chatWithContinue( + ctx context.Context, + modelID, prompt string, + cont openai.ContinuationFunc, + info *upstream.ServeInfo, + onFirstDelta func(), + emit func(text string) error, +) error { + return openai.RunContinued(ctx, prompt, cont, s.continueRounds(), + func(ctx context.Context, p string, emitText func(string) error) error { + return s.up.Chat(ctx, modelID, p, info, func(ev upstream.Event) error { + if ev.Kind == upstream.KindDelta && ev.Value != "" { + if onFirstDelta != nil { + onFirstDelta() + } + return emitText(ev.Value) + } + return nil + }) + }, + emit, + func(msg string, args ...any) { s.log.Info(msg, args...) }, + func(msg string, args ...any) { s.log.Warn(msg, args...) }, + ) +} diff --git a/internal/api/messages.go b/internal/api/messages.go index 4330183..95fc29a 100644 --- a/internal/api/messages.go +++ b/internal/api/messages.go @@ -66,15 +66,16 @@ func (s *Server) handleMessages(w http.ResponseWriter, r *http.Request) { s.met.IncInFlight(1) defer s.met.IncInFlight(-1) + cont := continuationFor(tools, messages, s.continueRounds()) if req.Stream { - s.streamMessages(ctx, w, req, prompt, site, tools) + s.streamMessages(ctx, w, req, prompt, site, tools, cont) return } - s.completeMessages(ctx, w, req, prompt, site, tools) + s.completeMessages(ctx, w, req, prompt, site, tools, cont) } -func (s *Server) completeMessages(ctx context.Context, w http.ResponseWriter, req anthropic.Request, prompt, site string, tools []openai.ToolDef) { - out := s.collect(ctx, req.Model, prompt, site) +func (s *Server) completeMessages(ctx context.Context, w http.ResponseWriter, req anthropic.Request, prompt, site string, tools []openai.ToolDef, cont openai.ContinuationFunc) { + out := s.collect(ctx, req.Model, prompt, site, cont) if out.err != nil { writeAnthropicUpstreamError(w, out.err) return @@ -99,7 +100,7 @@ func (s *Server) completeMessages(ctx context.Context, w http.ResponseWriter, re writeJSON(w, http.StatusOK, resp) } -func (s *Server) streamMessages(ctx context.Context, w http.ResponseWriter, req anthropic.Request, prompt, site string, tools []openai.ToolDef) { +func (s *Server) streamMessages(ctx context.Context, w http.ResponseWriter, req anthropic.Request, prompt, site string, tools []openai.ToolDef, cont openai.ContinuationFunc) { flusher, _ := w.(http.Flusher) w.Header().Set("Content-Type", "text/event-stream; charset=utf-8") w.Header().Set("Cache-Control", "no-cache") @@ -155,20 +156,18 @@ func (s *Server) streamMessages(ctx context.Context, w http.ResponseWriter, req var text strings.Builder var firstDelta time.Time - err := s.up.Chat(ctx, req.Model, prompt, info, func(ev upstream.Event) error { - if ev.Kind != upstream.KindDelta || ev.Value == "" { - return nil - } + err := s.chatWithContinue(ctx, req.Model, prompt, cont, info, func() { if firstDelta.IsZero() { firstDelta = time.Now() } - text.WriteString(ev.Value) + }, func(seg string) error { + text.WriteString(seg) if buffered { return nil } return writeEvent("content_block_delta", map[string]any{ "type": "content_block_delta", "index": 0, - "delta": map[string]any{"type": "text_delta", "text": ev.Value}, + "delta": map[string]any{"type": "text_delta", "text": seg}, }) }) diff --git a/internal/api/responses.go b/internal/api/responses.go index bd10ff3..f3e9836 100644 --- a/internal/api/responses.go +++ b/internal/api/responses.go @@ -58,11 +58,11 @@ func (s *Server) handleResponses(w http.ResponseWriter, r *http.Request) { defer cancel() if req.Stream { - s.streamResponses(ctx, w, req, prompt, site, tools) + s.streamResponses(ctx, w, req, prompt, site, tools, continuationFor(tools, messages, s.continueRounds())) return } - out := s.collect(ctx, req.Model, prompt, site) + out := s.collect(ctx, req.Model, prompt, site, continuationFor(tools, messages, s.continueRounds())) if out.err != nil { writeUpstreamError(w, out.err) return @@ -120,7 +120,7 @@ func responseOutputItems(text string, calls []openai.ToolCall) []map[string]any return output } -func (s *Server) streamResponses(ctx context.Context, w http.ResponseWriter, req openai.ResponsesRequest, prompt, site string, tools []openai.ToolDef) { +func (s *Server) streamResponses(ctx context.Context, w http.ResponseWriter, req openai.ResponsesRequest, prompt, site string, tools []openai.ToolDef, cont openai.ContinuationFunc) { flusher, _ := w.(http.Flusher) w.Header().Set("Content-Type", "text/event-stream; charset=utf-8") w.Header().Set("Cache-Control", "no-cache") @@ -158,18 +158,16 @@ func (s *Server) streamResponses(ctx context.Context, w http.ResponseWriter, req var text strings.Builder var firstDelta time.Time - err := s.up.Chat(ctx, req.Model, prompt, info, func(ev upstream.Event) error { - if ev.Kind != upstream.KindDelta || ev.Value == "" { - return nil - } + err := s.chatWithContinue(ctx, req.Model, prompt, cont, info, func() { if firstDelta.IsZero() { firstDelta = time.Now() } - text.WriteString(ev.Value) + }, func(seg string) error { + text.WriteString(seg) if buffered { return nil } - return write(map[string]any{"type": "response.output_text.delta", "delta": ev.Value}) + return write(map[string]any{"type": "response.output_text.delta", "delta": seg}) }) duration := time.Since(start) diff --git a/internal/config/config.go b/internal/config/config.go index da31e5a..f63d1f2 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -110,6 +110,7 @@ type Config struct { RetryBackoffSeconds float64 `toml:"retry_backoff_seconds" json:"retry_backoff_seconds"` UserAgent string `toml:"user_agent" json:"user_agent"` CrossSiteFailover bool `toml:"cross_site_failover" json:"cross_site_failover"` + ContinueRounds int `toml:"continue_rounds" json:"continue_rounds"` } `toml:"upstream" json:"upstream"` Runtime struct { @@ -198,6 +199,7 @@ func newConfig() *Config { c.Upstream.RetryBackoffSeconds = 1 c.Upstream.UserAgent = DefaultUserAgent c.Upstream.CrossSiteFailover = true + c.Upstream.ContinueRounds = 20 c.Runtime.DataDir = "./data" c.Turnstile = DefaultTurnstile() c.Sites = DefaultSites(c.Turnstile.SiteKey) @@ -297,6 +299,7 @@ func (c *Config) applyDefaults(md *toml.MetaData) { if md == nil { c.Turnstile.Enabled = true c.Upstream.AutoRefresh = true + c.Upstream.ContinueRounds = 20 for _, s := range c.Sites { s.Enabled = true } @@ -419,6 +422,9 @@ func (c *Config) Validate() error { if c.Upstream.StreamTimeout <= 0 { return errors.New("upstream.stream_timeout must be > 0") } + if c.Upstream.ContinueRounds < 0 { + c.Upstream.ContinueRounds = 0 + } if err := c.validateTurnstile(); err != nil { return err } diff --git a/internal/openai/continue.go b/internal/openai/continue.go new file mode 100644 index 0000000..1d5649e --- /dev/null +++ b/internal/openai/continue.go @@ -0,0 +1,246 @@ +package openai + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "unicode/utf8" +) + +// ContinuationContext returns a two-sided excerpt of partial: its first head +// runes keep the surrounding {"tool_calls": …} structure visible so the model +// never forgets it is mid-JSON, its last tail runes carry the resume point. +// Continuation turns embed only this excerpt: full replies grow past 20KB +// across rounds and the site degrades on huge prompts. +func ContinuationContext(partial string, head, tail int) string { + runes := []rune(partial) + if len(runes) <= head+tail { + return partial + } + mid := len(runes) - head - tail + return fmt.Sprintf("%s\n…（中间省略 %d 字符：以上是被截断回复的开头，以下是最新的末尾）…\n%s", + string(runes[:head]), mid, string(runes[len(runes)-tail:])) +} + +// ContinueInstruction builds the follow-up user turn that asks the model to +// resume a reply the channel cut off at its per-reply output cap. The site +// caps each generation (~1k tokens), so large tool-call JSON arrives in +// pieces; the model must continue byte-for-byte instead of restarting. +func ContinueInstruction(partial string, round, maxRounds int) string { + tail := partial + if n := len(tail); n > 80 { + tail = tail[n-80:] + } + deadline := fmt.Sprintf("这是第 %d/%d 次续写。", round, maxRounds) + if round >= maxRounds { + deadline += "这是最后一次续写机会：本轮必须先收束剩余内容（补齐未完成的必要结构），随后立即写收尾 8 字符序列闭合并结束回复，禁止再展开任何新段落。" + } else { + deadline += "每轮只能输出约 2300 字符；若主体已完成或接近完成，请立即闭合 JSON 收尾，不要无限扩展。" + } + return fmt.Sprintf( + "上一条 assistant 回复因通道输出上限被截断（上面保留了它的开头与最新末尾，中间已省略），%s目前已输出约 %d 个字符，最后 80 个字符是：\n%s\n\n"+ + "请只输出【续写片段】：从上面末尾片段的最后一个字符之后开始逐字符接着写，直到整个 JSON 完整闭合。"+ + "严格禁止：重复任何已输出内容、以 { 开头、重新生成 tool_calls 结构、任何解释或前后缀。"+ + "JSON 字符串内的换行必须写成两个字符的 \\n 转义序列、制表符写成 \\t，禁止输出真实的换行、制表或控制字符。"+ + "完成文件内容后必须立即闭合 JSON：依次写入 content 字符串的收尾引号、内层对象的 }、arguments 字符串的收尾引号、以及外层全部 } 与 ]，"+ + "文件写完后的收尾 8 个字符必须依次是：反斜杠、双引号、右花括号、双引号、右花括号、右花括号、右方括号、右花括号（即 \\\"}\"}}]}），一个都不能少，"+ + "然后立刻结束回复；禁止输出任何解释、Markdown 标记、使用说明、动画特点介绍或后记文字。"+ + "注意合理收尾：文件内容已经很长，若主体已经完成就立即闭合 JSON 结束，不要无限扩展新段落。"+ + "你的回复的第一个字符必须正好是被截断处的下一个字符。", + deadline, utf8.RuneCountInString(partial), tail) +} + +// ToolCallTruncated reports whether text looks like a forced tool-call JSON +// reply (see toolsPrompt) that the channel cut off before it could close. +// Fenced ```json blocks are unwrapped first, mirroring TryParseToolCalls. +func ToolCallTruncated(text string) bool { + cand := strings.TrimSpace(text) + if !strings.HasPrefix(cand, "{") { + if strings.Contains(cand, "```json") { + cand = strings.TrimSpace(firstBlock(strings.Split(cand, "```json")[1])) + } else if strings.Contains(cand, "```") { + cand = strings.TrimSpace(firstBlock(strings.Split(cand, "```")[1])) + } + } + if !strings.HasPrefix(cand, `{"tool_calls"`) { + return false + } + return !json.Valid([]byte(cand)) +} + +// EscapeControlChars rewrites raw control characters that appear inside JSON +// string literals as \n / \r / \t escapes. Whitespace outside string literals +// (pretty-printed structure) is left untouched. It equals +// EscapeControlsAfter("", s) and suits complete documents. +func EscapeControlChars(s string) string { + return EscapeControlsAfter("", s) +} + +// EscapeControlsAfter escapes control characters in s while assuming s +// continues prefix inside the same JSON document: the string-literal state at +// the end of prefix carries over, which matters because continuation turns +// resume mid-string. +func EscapeControlsAfter(prefix, s string) string { + if !needsEscape(s) { + return s + } + st := escapeState{} + if prefix != "" { + var throw strings.Builder + st.feed(prefix, &throw) + } + var out strings.Builder + out.Grow(len(s) + 16) + st.feed(s, &out) + return out.String() +} + +type escapeState struct { + inString bool + escaped bool +} + +func (st *escapeState) feed(s string, out *strings.Builder) { + for i := 0; i < len(s); i++ { + c := s[i] + if !st.inString { + if c == '"' { + st.inString = true + } + out.WriteByte(c) + continue + } + if st.escaped { + st.escaped = false + out.WriteByte(c) + continue + } + switch c { + case '\\': + st.escaped = true + out.WriteByte(c) + case '"': + st.inString = false + out.WriteByte(c) + case '\n': + out.WriteString(`\n`) + case '\r': + out.WriteString(`\r`) + case '\t': + out.WriteString(`\t`) + default: + if c < 0x20 { + fmt.Fprintf(out, `\u%04x`, c) + } else { + out.WriteByte(c) + } + } + } +} + +func needsEscape(s string) bool { + for i := 0; i < len(s); i++ { + if s[i] < 0x20 { + return true + } + } + return false +} + +// TrimOverlap drops the longest prefix of head that duplicates the tail of +// acc: continuation turns sometimes re-emit the last few characters of the +// truncated reply before carrying on. +func TrimOverlap(acc, head string) string { + max := len(acc) + if len(head) < max { + max = len(head) + } + if max > 64 { + max = 64 + } + for k := max; k > 0; k-- { + if strings.HasSuffix(acc, head[:k]) { + return head[k:] + } + } + return head +} + +// SendFunc runs one upstream conversation for prompt and streams its text +// through emit. Emit errors abort the conversation. +type SendFunc func(ctx context.Context, prompt string, emit func(text string) error) error + +// ContinuationFunc builds the follow-up prompt for partial. round is 1-based +// and lets the instruction tell the model how many attempts it has left. +type ContinuationFunc func(partial string, round int) string + +// RunContinued sends first through send, then — while the accumulated text is +// a truncated tool-call JSON — keeps sending cont(partial) turns to resume it +// until the JSON closes or rounds run out. Every stitched piece reaches emit +// in order, so callers see one continuous reply. Continuation failures only +// warn: the caller falls back to whatever text arrived, exactly as before. +func RunContinued( + ctx context.Context, + first string, + cont ContinuationFunc, + rounds int, + send SendFunc, + emit func(text string) error, + logf func(msg string, args ...any), + warn func(msg string, args ...any), +) error { + var full strings.Builder + err := send(ctx, first, func(text string) error { + full.WriteString(text) + return emit(text) + }) + if err != nil { + return err + } + if cont == nil || rounds <= 0 { + return nil + } + for round := 0; round < rounds && ToolCallTruncated(full.String()); round++ { + if ctx.Err() != nil { + return ctx.Err() + } + partial := full.String() + next := cont(partial, round+1) + var seg strings.Builder + var sendErr error + for attempt := 0; attempt < 3; attempt++ { + seg.Reset() + sendErr = send(ctx, next, func(text string) error { + seg.WriteString(text) + return nil + }) + if sendErr == nil { + break + } + if ctx.Err() != nil { + return ctx.Err() + } + warn("tool-call continuation turn retrying", "round", round+1, "attempt", attempt+1, "error", sendErr) + } + if sendErr != nil { + warn("tool-call continuation turn failed", "round", round+1, "prompt_chars", len(next), "error", sendErr) + break + } + raw := seg.String() + piece := TrimOverlap(partial, EscapeControlsAfter(partial, raw)) + if piece == "" { + warn("tool-call continuation produced no new text", "round", round+1, "seg_chars", len(raw)) + break + } + full.WriteString(piece) + logf("tool-call continuation stitched", "round", round+1, "added_chars", len(piece), "seg_chars", len(raw)) + if err := emit(piece); err != nil { + return err + } + } + if ToolCallTruncated(full.String()) { + warn("tool-call continuation exhausted — returning truncated reply as-is") + } + return nil +} diff --git a/internal/openai/continue_test.go b/internal/openai/continue_test.go new file mode 100644 index 0000000..fc385db --- /dev/null +++ b/internal/openai/continue_test.go @@ -0,0 +1,337 @@ +package openai + +import ( + "context" + "encoding/json" + "errors" + "strings" + "testing" +) + +func TestToolCallTruncated(t *testing.T) { + valid := `{"tool_calls": [{"id": "a", "type": "function", "function": {"name": "write", "arguments": "{\"filePath\": \"x\"}"}}]}` + truncated := valid[:40] + prose := "the model mentions tool_calls but writes no JSON at all" + fenced := "```json\n" + truncated + + cases := []struct { + in string + want bool + }{ + {valid, false}, + {truncated, true}, + {prose, false}, + {fenced, true}, + {"", false}, + {`{"other": true}`, false}, + } + for _, c := range cases { + if got := ToolCallTruncated(c.in); got != c.want { + t.Errorf("ToolCallTruncated(%q) = %v, want %v", c.in, got, c.want) + } + } +} + +func TestEscapeControlChars(t *testing.T) { + cases := []struct{ in, want string }{ + {`{"a": "line1` + "\n" + `line2"}`, `{"a": "line1\nline2"}`}, + {`{"a": "tab` + "\t" + `x"}`, `{"a": "tab\tx"}`}, + // Escaped sequences inside strings stay as-is. + {`{"a": "already\nescaped"}`, `{"a": "already\nescaped"}`}, + // Pretty-printed structure newlines outside strings survive. + {"{\n \"a\": \"x\"\n}", "{\n \"a\": \"x\"\n}"}, + // Real newline inside a second string while first had escapes. + {`{"a": "x\ny", "b": "p` + "\n" + `q"}`, `{"a": "x\ny", "b": "p\nq"}`}, + // Other control characters get \u-escaped. + {`{"a": "` + string(rune(0x01)) + `"}`, `{"a": "\u0001"}`}, + } + for _, c := range cases { + if got := EscapeControlChars(c.in); got != c.want { + t.Errorf("EscapeControlChars(%q) = %q, want %q", c.in, got, c.want) + } + } + // Round-trip: repaired text must be valid JSON with the same value. + raw := "{\"a\": \"line1\nline2\"}" + fixed := EscapeControlChars(raw) + var v struct { + A string `json:"a"` + } + if err := json.Unmarshal([]byte(fixed), &v); err != nil { + t.Fatalf("repaired JSON invalid: %v", err) + } + if v.A != "line1\nline2" { + t.Fatalf("repaired value = %q", v.A) + } +} + +func TestTrimOverlap(t *testing.T) { + cases := []struct{ acc, head, want string }{ + {"...fillStyle='#F4", "F4E2B8';draw()", "E2B8';draw()"}, + {"...fillStyle='#F4", "E2B8';draw()", "E2B8';draw()"}, + {"abcdef", "xyz", "xyz"}, + {"abcabcabc", "abcabcTAIL", "TAIL"}, + {"", "anything", "anything"}, + } + for _, c := range cases { + if got := TrimOverlap(c.acc, c.head); got != c.want { + t.Errorf("TrimOverlap(%q, %q) = %q, want %q", c.acc, c.head, got, c.want) + } + } +} + +func TestEscapeControlsAfterCarriesPrefixState(t *testing.T) { + prefix := `{"tool_calls": [{"function": {"arguments": "{\"content\": \"abc` + head := "def\n" + `ghi\""}}]}` + got := EscapeControlsAfter(prefix, head) + want := `def\nghi\""}}]}` + if got != want { + t.Fatalf("EscapeControlsAfter = %q, want %q", got, want) + } + // Without prefix context the same head starts outside a string, so the + // newline is structure whitespace and must survive. + if plain := EscapeControlChars(head); plain != head { + t.Fatalf("EscapeControlChars = %q, want unchanged", plain) + } + full := prefix + got + if !json.Valid([]byte(full)) { + t.Fatalf("repaired document invalid: %s", full) + } +} + +func TestContinueInstructionMentionsTail(t *testing.T) { + partial := strings.Repeat("x", 200) + "TAIL80CHARS" + got := ContinueInstruction(partial, 3, 12) + if !strings.Contains(got, "3/12") { + t.Errorf("instruction missing round progress: %s", got) + } + if !strings.Contains(got, "TAIL80CHARS") { + t.Errorf("instruction missing tail: %s", got) + } + if strings.Contains(got, strings.Repeat("x", 100)) { + t.Errorf("instruction embedded too much of the partial") + } + if !strings.Contains(got, `\n`) { + t.Errorf("instruction missing escape-sequence rule") + } +} + +func TestBalanceJSON(t *testing.T) { + cases := []struct{ name, in, want string }{ + {"unterminated content", `{"filePath": "x.html", "content": "`, `{"filePath": "x.html", "content": ""}`}, + {"trailing comma", `{"a": "b",`, `{"a": "b"}`}, + {"trailing colon", `{"a":`, `{"a": ""}`}, + {"dangling escape", `{"a": "b\`, `{"a": "b"}`}, + {"nested array", `{"a": [1, 2`, `{"a": [1, 2]}`}, + {"already valid", `{"a": "b"}`, `{"a": "b"}`}, + {"comma inside string survives", `{"a": "b,c`, `{"a": "b,c"` + `}`}, + } + for _, c := range cases { + if got := balanceJSON(c.in); got != c.want { + t.Errorf("%s: balanceJSON(%q) = %q, want %q", c.name, c.in, got, c.want) + } + } + // Unfixable damage passes through untouched. + broken := `{"a": 12.}` + if got := balanceJSON(broken); got != broken { + t.Errorf("unfixable: got %q", got) + } +} + +func TestTryParseToolCallsHealsBrokenInnerJSON(t *testing.T) { + text := `{"tool_calls": [{"id": "a", "type": "function", "function": {"name": "write", "arguments": "{\"content\": \""}}]}` + _, calls, ok := TryParseToolCalls(text) + if !ok { + t.Fatal("expected healed tool call") + } + var inner struct { + Content string `json:"content"` + } + if err := json.Unmarshal([]byte(calls[0].Function.Arguments), &inner); err != nil { + t.Fatalf("healed arguments invalid: %v", err) + } + if inner.Content != "" { + t.Fatalf("content = %q", inner.Content) + } +} + +func TestTryParseToolCallsRejectsUnfixableInnerJSON(t *testing.T) { + text := `{"tool_calls": [{"id": "a", "type": "function", "function": {"name": "write", "arguments": "{\"content\": 12.}"}}]}` + _, calls, ok := TryParseToolCalls(text) + if ok { + t.Fatalf("expected rejection, got calls=%v", calls) + } +} + +func TestRunContinuedStitches(t *testing.T) { + turns := []string{ + `{"tool_calls": [{"id": "a", "function": {"name": "write", "arguments": "{\"content\": \"part1`, + // Continuation repeats one char, emits a raw newline mid-string, then + // closes the inner JSON, the arguments string and the outer document. + "1\n" + `part2\"}"}}]}`, + } + prompts := 0 + var got strings.Builder + err := RunContinued(context.Background(), + "first-prompt", + func(partial string, _ int) string { + if !strings.Contains(partial, "part1") { + t.Errorf("cont prompt built from unexpected partial: %s", partial) + } + return "continue-prompt" + }, + 3, + func(_ context.Context, prompt string, emit func(string) error) error { + prompts++ + idx := prompts - 1 + if idx >= len(turns) { + t.Fatalf("unexpected extra send #%d", prompts) + } + if idx > 0 && prompt != "continue-prompt" { + t.Errorf("round %d prompt = %q", idx, prompt) + } + return emit(turns[idx]) + }, + func(text string) error { got.WriteString(text); return nil }, + func(string, ...any) {}, + func(string, ...any) {}, + ) + if err != nil { + t.Fatalf("RunContinued: %v", err) + } + if prompts != 2 { + t.Fatalf("prompts = %d, want 2", prompts) + } + want := `{"tool_calls": [{"id": "a", "function": {"name": "write", "arguments": "{\"content\": \"part1\npart2\"}"}}]}` + if got.String() != want { + t.Fatalf("stitched = %q\nwant %q", got.String(), want) + } + if ToolCallTruncated(got.String()) { + t.Fatal("stitched reply still truncated") + } + _, calls, ok := TryParseToolCalls(got.String()) + if !ok { + t.Fatalf("TryParseToolCalls failed on stitched reply") + } + var inner struct { + Content string `json:"content"` + } + if err := json.Unmarshal([]byte(calls[0].Function.Arguments), &inner); err != nil { + t.Fatalf("arguments JSON invalid: %v", err) + } + if inner.Content != "part1\npart2" { + t.Fatalf("content = %q", inner.Content) + } +} + +func TestRunContinuedPassthroughWithoutCont(t *testing.T) { + prompts := 0 + var got strings.Builder + err := RunContinued(context.Background(), "p", nil, 6, + func(_ context.Context, prompt string, emit func(string) error) error { + prompts++ + return emit(`{"tool_calls": [{"trunc`) + }, + func(text string) error { got.WriteString(text); return nil }, + func(string, ...any) {}, func(string, ...any) {}, + ) + if err != nil { + t.Fatalf("RunContinued: %v", err) + } + if prompts != 1 { + t.Fatalf("prompts = %d, want 1", prompts) + } + if !ToolCallTruncated(got.String()) { + t.Fatal("passthrough should keep the truncated reply") + } +} + +func TestRunContinuedContinuationFailureKeepsPartial(t *testing.T) { + prompts := 0 + warns := 0 + var got strings.Builder + err := RunContinued(context.Background(), "first", + func(string, int) string { return "next" }, 3, + func(_ context.Context, _ string, emit func(string) error) error { + prompts++ + if prompts == 1 { + return emit(`{"tool_calls": [{"trunc`) + } + return errors.New("boom") + }, + func(text string) error { got.WriteString(text); return nil }, + func(string, ...any) {}, + func(string, ...any) { warns++ }, + ) + if err != nil { + t.Fatalf("RunContinued returned error despite partial text: %v", err) + } + if warns != 5 { + t.Fatalf("warns = %d, want 5 (3 retry notices + turn failure + exhaustion)", warns) + } + if got.String() != `{"tool_calls": [{"trunc` { + t.Fatalf("got = %q", got.String()) + } +} + +func TestRunContinuedStopsWhenClosed(t *testing.T) { + prompts := 0 + err := RunContinued(context.Background(), "first", + func(string, int) string { return "next" }, 6, + func(_ context.Context, _ string, emit func(string) error) error { + prompts++ + return emit(`{"tool_calls": [{"id": "a", "type": "function", "function": {"name": "f", "arguments": "{}"}}]}`) + }, + func(string) error { return nil }, + func(string, ...any) {}, func(string, ...any) {}, + ) + if err != nil { + t.Fatalf("RunContinued: %v", err) + } + if prompts != 1 { + t.Fatalf("prompts = %d, want 1 (valid JSON must not trigger continuation)", prompts) + } +} + +func TestRunContinuedEmitErrorPropagates(t *testing.T) { + sentinel := errors.New("emit failed") + err := RunContinued(context.Background(), "first", nil, 0, + func(_ context.Context, _ string, emit func(string) error) error { + return emit("text") + }, + func(string) error { return sentinel }, + func(string, ...any) {}, func(string, ...any) {}, + ) + if !errors.Is(err, sentinel) { + t.Fatalf("err = %v, want sentinel", err) + } +} + +func TestTryParseToolCallsShortCloseSequence(t *testing.T) { + inner := `{\"filePath\": \"a.html\", \"content\": \"\"}` + raw := `{"tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "write", "arguments": "` + inner + `}}]}` + + remaining, calls, ok := TryParseToolCalls(raw) + if !ok { + t.Fatal("TryParseToolCalls ok = false, want true (short close sequence must heal)") + } + if remaining != "" { + t.Fatalf("remaining = %q, want empty", remaining) + } + if len(calls) != 1 { + t.Fatalf("calls = %d, want 1", len(calls)) + } + args := calls[0].Function.Arguments + if !json.Valid([]byte(args)) { + t.Fatalf("args invalid: %s", args) + } + var obj struct { + FilePath string `json:"filePath"` + Content string `json:"content"` + } + if err := json.Unmarshal([]byte(args), &obj); err != nil { + t.Fatal(err) + } + if obj.Content != "" { + t.Fatalf("content = %q, want html", obj.Content) + } +} diff --git a/internal/openai/toolcalls.go b/internal/openai/toolcalls.go index 8885980..5bd37f8 100644 --- a/internal/openai/toolcalls.go +++ b/internal/openai/toolcalls.go @@ -42,13 +42,14 @@ func TryParseToolCalls(text string) (string, []ToolCall, bool) { Function json.RawMessage `json:"function"` } `json:"tool_calls"` } - if err := json.Unmarshal([]byte(cand), &payload); err != nil { + if !decodeToolPayload(cand, &payload) { continue } if payload.ToolCalls == nil { continue } calls := make([]ToolCall, 0, len(payload.ToolCalls)) + bad := false for i, tc := range payload.ToolCalls { var fn struct { Name string `json:"name"` @@ -59,7 +60,25 @@ func TryParseToolCalls(text string) (string, []ToolCall, bool) { if len(fn.Arguments) > 0 { var s string if err := json.Unmarshal(fn.Arguments, &s); err == nil { - args = s + // Continuation stitching may have introduced raw control + // characters into the inner JSON; heal them and close any + // trailing string/brace the model forgot before the + // channel cap let it stop. + s = EscapeControlChars(s) + args = balanceJSON(s) + if !json.Valid([]byte(args)) { + // The model's close sequence may be one brace or + // quote short, leaving stray closers after the inner + // object: keep the first complete JSON value and + // drop the dangling tail. + if prefix, ok := jsonPrefix(s); ok { + args = prefix + } + } + if !json.Valid([]byte(args)) { + bad = true + break + } } else { // arguments already delivered as an object var obj any @@ -87,6 +106,9 @@ func TryParseToolCalls(text string) (string, []ToolCall, bool) { }, }) } + if bad { + continue + } remaining := strings.ReplaceAll(text, cand, "") remaining = strings.ReplaceAll(remaining, "```json", "") remaining = strings.ReplaceAll(remaining, "```", "") @@ -95,6 +117,37 @@ func TryParseToolCalls(text string) (string, []ToolCall, bool) { return text, nil, false } +// decodeToolPayload unmarshals cand, falling back to suffix balancing (the +// model broke out mid-string) and to the first complete JSON value (stray +// closers after a complete object). +func decodeToolPayload(cand string, payload any) bool { + if err := json.Unmarshal([]byte(cand), payload); err == nil { + return true + } + if balanced := balanceJSON(cand); balanced != cand { + if err := json.Unmarshal([]byte(balanced), payload); err == nil { + return true + } + } + if prefix, ok := jsonPrefix(cand); ok { + if err := json.Unmarshal([]byte(prefix), payload); err == nil { + return true + } + } + return false +} + +// jsonPrefix returns the first complete JSON value in s, ignoring any +// trailing garbage after it. +func jsonPrefix(s string) (string, bool) { + dec := json.NewDecoder(strings.NewReader(s)) + var raw json.RawMessage + if err := dec.Decode(&raw); err != nil || len(raw) == 0 { + return "", false + } + return string(raw), true +} + func firstBlock(s string) string { if i := strings.Index(s, "```"); i >= 0 { return s[:i] @@ -102,6 +155,90 @@ func firstBlock(s string) string { return s } +// balanceJSON repairs an arguments string the model closed prematurely: a +// trailing unterminated string, a dangling comma/colon or unclosed braces +// (the channel cap sometimes lets a reply stop right before its inner JSON +// closes). It returns the original text untouched when the damage is not a +// simple missing-suffix problem. +func balanceJSON(s string) string { + if json.Valid([]byte(s)) { + return s + } + var stack []byte + inString, escaped := false, false + for i := 0; i < len(s); i++ { + c := s[i] + if !inString { + switch c { + case '"': + inString = true + case '{', '[': + stack = append(stack, c) + case '}', ']': + if len(stack) > 0 { + stack = stack[:len(stack)-1] + } + } + continue + } + if escaped { + escaped = false + continue + } + switch c { + case '\\': + escaped = true + case '"': + inString = false + } + } + + out := s + // A dangling escape (…\) is a partial escape sequence: drop it. + if inString && escaped { + out = out[:len(out)-1] + escaped = false + } + if !inString { + out = strings.TrimRight(out, " \t\r\n") + addedValue := false + for len(out) > 0 { + last := out[len(out)-1] + if last == ',' || last == ':' { + if last == ':' { + addedValue = true + } + out = strings.TrimRight(out[:len(out)-1], " \t\r\n") + continue + } + break + } + var trimmed strings.Builder + trimmed.WriteString(out) + if addedValue { + trimmed.WriteString(": \"\"") + } + out = trimmed.String() + } + var b strings.Builder + b.Grow(len(out) + len(stack) + 4) + b.WriteString(out) + if inString { + b.WriteByte('"') + } + for i := len(stack) - 1; i >= 0; i-- { + if stack[i] == '{' { + b.WriteByte('}') + } else { + b.WriteByte(']') + } + } + if json.Valid([]byte(b.String())) { + return b.String() + } + return s +} + func randomHex(n int) string { b := make([]byte, n) if _, err := rand.Read(b); err != nil {