From fe1d5b534882f22a6366d71564bac830cf9cc27b Mon Sep 17 00:00:00 2001 From: lutc5 Date: Wed, 6 May 2026 15:58:37 +0800 Subject: [PATCH] Fix tool loop handling and count tokens endpoint --- desktop/app.go | 15 ++++- desktop/frontend/src/App.vue | 8 ++- internal/httpapi/server.go | 65 ++++++++++++++++++++ internal/httpapi/server_test.go | 58 +++++++++++++++++ internal/toolemulation/toolemulation.go | 21 ++++++- internal/toolemulation/toolemulation_test.go | 22 +++++++ 6 files changed, 182 insertions(+), 7 deletions(-) diff --git a/desktop/app.go b/desktop/app.go index 6e446c7..e888502 100644 --- a/desktop/app.go +++ b/desktop/app.go @@ -231,8 +231,19 @@ func (a *App) forceQuit() { a.mu.Unlock() a.emitLog("info", "正在停止代理并退出应用") - if err := a.StopProxy(); err != nil { - runtime.LogWarningf(a.ctx, "stop proxy before exit failed: %v", err) + + done := make(chan struct{}) + go func() { + if err := a.StopProxy(); err != nil { + runtime.LogWarningf(a.ctx, "stop proxy before exit failed: %v", err) + } + close(done) + }() + + select { + case <-done: + case <-time.After(1200 * time.Millisecond): + runtime.LogWarning(a.ctx, "force quit continuing before proxy shutdown completed") } os.Exit(0) } diff --git a/desktop/frontend/src/App.vue b/desktop/frontend/src/App.vue index 68060b0..8dc8ca1 100644 --- a/desktop/frontend/src/App.vue +++ b/desktop/frontend/src/App.vue @@ -15,6 +15,7 @@ const status = ref({ running: false, addr: '', models: 0 }) const toast = ref('') const themeMode = ref(localStorage.getItem('lingma-theme-mode') || 'system') const appliedTheme = ref('light') +const forceQuitting = ref(false) let systemThemeQuery = null let toastTimer = null @@ -106,12 +107,13 @@ async function copyEndpoint() { } async function forceQuitApp() { - const confirmed = window.confirm('确定要停止代理并退出应用吗？') - if (!confirmed) return + if (forceQuitting.value) return + forceQuitting.value = true showToast('正在停止代理并退出应用...') try { await ForceQuitApp() } catch (e) { + forceQuitting.value = false addLog('error', '退出应用失败：' + (e.message || String(e))) } } @@ -271,7 +273,7 @@ onUnmounted(() => { - diff --git a/internal/httpapi/server.go b/internal/httpapi/server.go index 778ca6e..482b733 100644 --- a/internal/httpapi/server.go +++ b/internal/httpapi/server.go @@ -117,6 +117,7 @@ func NewServer(addr string, svc *service.Service) *Server { mux.HandleFunc("/v1/props", s.handleModelProps) mux.HandleFunc("/props", s.handleModelProps) mux.HandleFunc("/version", s.handleVersion) + mux.HandleFunc("/v1/messages/count_tokens", s.handleAnthropicCountTokens) mux.HandleFunc("/v1/messages", s.handleAnthropicMessages) mux.HandleFunc("/v1/chat/completions", s.handleOpenAIChatCompletions) mux.HandleFunc("/api/v1/chat/completions", s.handleOpenAIChatCompletions) @@ -446,6 +447,27 @@ func (s *Server) handleVersion(w http.ResponseWriter, r *http.Request) { }) } +func (s *Server) handleAnthropicCountTokens(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodOptions { + w.WriteHeader(http.StatusNoContent) + return + } + if r.Method != http.MethodPost { + writeAnthropicError(w, http.StatusMethodNotAllowed, "invalid_request_error", "method not allowed") + return + } + + var req anthropicRequest + if err := decodeJSON(r, &req); err != nil { + writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", err.Error()) + return + } + + writeJSON(w, http.StatusOK, map[string]any{ + "input_tokens": estimateAnthropicInputTokens(req), + }) +} + func (s *Server) handleAnthropicMessages(w http.ResponseWriter, r *http.Request) { if r.Method == http.MethodOptions { w.WriteHeader(http.StatusNoContent) @@ -1292,6 +1314,9 @@ func anthropicHostedWebSearchCall(req anthropicRequest) (toolemulation.ToolCall, if !hasAnthropicHostedWebSearchTool(req.Tools) { return toolemulation.ToolCall{}, false } + if hasAnthropicToolResult(req.Messages) { + return toolemulation.ToolCall{}, false + } if !anthropicHostedWebSearchRequested(req.Tools, req.ToolChoice) { return toolemulation.ToolCall{}, false } @@ -1307,6 +1332,46 @@ func anthropicHostedWebSearchCall(req anthropicRequest) (toolemulation.ToolCall, }, true } +func hasAnthropicToolResult(messages []rawMessage) bool { + for _, message := range messages { + items, ok := message.Content.([]any) + if !ok { + continue + } + for _, item := range items { + m, ok := item.(map[string]any) + if ok && stringFromAny(m["type"]) == "tool_result" { + return true + } + } + } + return false +} + +func estimateAnthropicInputTokens(req anthropicRequest) int { + payload := map[string]any{ + "model": req.Model, + "system": req.System, + "messages": req.Messages, + "tools": req.Tools, + "tool_choice": req.ToolChoice, + "thinking": req.Thinking, + } + raw, err := json.Marshal(payload) + if err != nil { + return 1 + } + runes := len([]rune(string(raw))) + if runes == 0 { + return 1 + } + tokens := (runes + 2) / 3 + if tokens < 1 { + return 1 + } + return tokens +} + func hasAnthropicHostedWebSearchTool(raw any) bool { items, ok := raw.([]any) if !ok { diff --git a/internal/httpapi/server_test.go b/internal/httpapi/server_test.go index b302bf3..efc60fd 100644 --- a/internal/httpapi/server_test.go +++ b/internal/httpapi/server_test.go @@ -218,6 +218,64 @@ func TestAnthropicHostedWebSearchCallIgnoresRegularClientWebSearch(t *testing.T) } } +func TestAnthropicHostedWebSearchCallIgnoresToolResultFollowup(t *testing.T) { + req := anthropicRequest{ + Tools: []any{ + map[string]any{ + "name": "web_search", + "type": "web_search_20250305", + }, + }, + ToolChoice: map[string]any{ + "type": "tool", + "name": "web_search", + }, + Messages: []rawMessage{{ + Role: "user", + Content: []any{ + map[string]any{ + "type": "tool_result", + "tool_use_id": "toolu_123", + "content": "result", + }, + }, + }}, + } + + if _, ok := anthropicHostedWebSearchCall(req); ok { + t.Fatal("hosted web_search should not short-circuit after a tool_result") + } +} + +func TestAnthropicCountTokensEndpoint(t *testing.T) { + server := NewServer("", service.New(service.Config{ + Model: "Qwen3-Coder", + Timeout: time.Second, + })) + + req := httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", strings.NewReader(`{ + "model":"kmodel", + "max_tokens":128, + "system":"You are concise.", + "messages":[{"role":"user","content":"hello"}], + "tools":[{"name":"read_file","input_schema":{"type":"object","properties":{"file_path":{"type":"string"}},"required":["file_path"]}}] + }`)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + server.http.Handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d body = %s", rec.Code, rec.Body.String()) + } + var body map[string]any + if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil { + t.Fatal(err) + } + if body["input_tokens"].(float64) <= 0 { + t.Fatalf("input_tokens = %#v", body["input_tokens"]) + } +} + func TestDiscoveryCompatibilityEndpoints(t *testing.T) { server := NewServer("", service.New(service.Config{ Model: "Qwen3-Coder", diff --git a/internal/toolemulation/toolemulation.go b/internal/toolemulation/toolemulation.go index 602fc80..8b62a60 100644 --- a/internal/toolemulation/toolemulation.go +++ b/internal/toolemulation/toolemulation.go @@ -283,10 +283,11 @@ func ActionOutputPrompt(toolCallID string, output string) string { if output == "" { return "" } + next := "Based on the tool result above, answer the user's request directly if you have enough information. Only use another structured action block if a specific missing fact still requires another tool call." if id := strings.TrimSpace(toolCallID); id != "" { - return "Tool result for " + id + ":\n" + output + "\n\nBased on the tool result above, continue with the next appropriate action using the structured format." + return "Tool result for " + id + ":\n" + output + "\n\n" + next } - return "Tool result:\n" + output + "\n\nBased on the tool result above, continue with the next appropriate action using the structured format." + return "Tool result:\n" + output + "\n\n" + next } func ActionBlockExample(tools []ToolDef) string { @@ -629,6 +630,9 @@ func ParseActionBlocks(text string, tools []ToolDef, cfg Config) ([]ToolCall, st // Filter arguments against the tool's input schema to strip unknown params if schema, ok := toolSchemaMap[call.Name]; ok && len(schema) > 0 { call.Arguments = filterArgsBySchema(call.Arguments, schema) + if !hasRequiredArgs(call.Arguments, schema) { + continue + } } calls = append(calls, call) spans = append(spans, span{start: start, end: end + 3}) @@ -1011,6 +1015,19 @@ func filterArgsBySchema(args map[string]any, schema map[string]any) map[string]a return out } +func hasRequiredArgs(args map[string]any, schema map[string]any) bool { + for _, key := range requiredKeys(schema) { + value, ok := args[key] + if !ok { + return false + } + if s, ok := value.(string); ok && strings.TrimSpace(s) == "" { + return false + } + } + return true +} + func cloneMap(src map[string]any) map[string]any { if src == nil { return nil diff --git a/internal/toolemulation/toolemulation_test.go b/internal/toolemulation/toolemulation_test.go index e48f83c..6c8119b 100644 --- a/internal/toolemulation/toolemulation_test.go +++ b/internal/toolemulation/toolemulation_test.go @@ -154,3 +154,25 @@ func TestParseActionBlocksMapsReadAlias(t *testing.T) { t.Fatalf("calls = %+v", calls) } } + +func TestParseActionBlocksDropsCallsMissingRequiredArgs(t *testing.T) { + text := "```json action\n{\"tool\":\"Read\",\"parameters\":{\"path\":\"/tmp/a.txt\"}}\n```" + calls, clean, err := ParseActionBlocks(text, []ToolDef{{ + Name: "Read", + InputSchema: map[string]any{ + "properties": map[string]any{ + "file_path": map[string]any{"type": "string"}, + }, + "required": []any{"file_path"}, + }, + }}, Config{}) + if err != nil { + t.Fatal(err) + } + if len(calls) != 0 { + t.Fatalf("expected no calls, got %+v", calls) + } + if !strings.Contains(clean, "\"path\"") { + t.Fatalf("clean should preserve unparseable action block, got %q", clean) + } +}