diff --git a/internal/exec/litellm.go b/internal/exec/litellm.go index 278f8eb..ab22a55 100644 --- a/internal/exec/litellm.go +++ b/internal/exec/litellm.go @@ -6,6 +6,7 @@ import ( "encoding/json" "fmt" "net/http" + "strings" "time" ) @@ -93,5 +94,34 @@ func (e *LiteLLMExecutor) Complete(ctx context.Context, model, system, user stri return "", 0, fmt.Errorf("litellm: no choices in response") } - return chatResp.Choices[0].Message.Content, durationMs, nil + return stripResultJSON(chatResp.Choices[0].Message.Content), durationMs, nil +} + +// stripResultJSON removes trailing ```json blocks that match the old structured +// result schema (containing "status" and "phase" keys). Some local models produce +// correct markdown prose but then append the old JSON format out of habit. +func stripResultJSON(text string) string { + const fence = "```json" + idx := len(text) - 1 + // Walk backwards past trailing whitespace. + for idx >= 0 && (text[idx] == '\n' || text[idx] == '\r' || text[idx] == ' ') { + idx-- + } + // Must end with closing fence. + if idx < 2 || text[idx-2:idx+1] != "```" { + return text + } + // Find the matching opening fence. + start := len(text[:idx-2]) - 1 + for start >= 0 { + if start+len(fence) <= len(text) && text[start:start+len(fence)] == fence { + block := text[start : idx+1] + if strings.Contains(block, `"status"`) && strings.Contains(block, `"phase"`) { + return strings.TrimRight(text[:start], " \t\r\n") + } + break + } + start-- + } + return text } diff --git a/internal/exec/litellm_test.go b/internal/exec/litellm_test.go index 71afc96..9cdf944 100644 --- a/internal/exec/litellm_test.go +++ b/internal/exec/litellm_test.go @@ -80,6 +80,40 @@ func TestLiteLLMErrorOnEmptyChoices(t *testing.T) { assert.ErrorContains(t, err, "no choices") } +func TestLiteLLMStripsTrailingResultJSON(t *testing.T) { + content := "## Hypotheses\n\n**H1 (high):** nil map access.\n\n```json\n{\n \"status\": \"pass\",\n \"phase\": \"debug\",\n \"skill\": \"debug\"\n}\n```" + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(chatResponse(t, content)) + })) + defer srv.Close() + + ex := iexec.NewLiteLLM(srv.URL, "", 5*time.Second) + text, _, err := ex.Complete(context.Background(), "model", "sys", "user") + require.NoError(t, err) + assert.Contains(t, text, "nil map access") + assert.NotContains(t, text, `"status"`) + assert.NotContains(t, text, "```json") +} + +func TestLiteLLMKeepsNonResultJSONFence(t *testing.T) { + // A json block that is part of the actual answer (no status/phase) should be kept. + content := "Use this config:\n\n```json\n{\"model\": \"koala/phi4\"}\n```" + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(chatResponse(t, content)) + })) + defer srv.Close() + + ex := iexec.NewLiteLLM(srv.URL, "", 5*time.Second) + text, _, err := ex.Complete(context.Background(), "model", "sys", "user") + require.NoError(t, err) + assert.Contains(t, text, `"model"`) + assert.Contains(t, text, "```json") +} + func TestLiteLLMRespectsContextCancellation(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel()