1  package tools
  2  
  3  import (
  4  	"context"
  5  	"encoding/json"
  6  	"fmt"
  7  	"strings"
  8  
  9  	"github.com/Kocoro-lab/ShanClaw/internal/agent"
 10  	"github.com/Kocoro-lab/ShanClaw/internal/client"
 11  )
 12  
 13  // ServerTool wraps a server-side tool schema and proxies execution
 14  // through the gateway's /api/v1/tools/{name}/execute endpoint.
 15  type ServerTool struct {
 16  	schema  client.ServerToolSchema
 17  	gateway *client.GatewayClient
 18  }
 19  
 20  func NewServerTool(schema client.ServerToolSchema, gateway *client.GatewayClient) *ServerTool {
 21  	return &ServerTool{schema: schema, gateway: gateway}
 22  }
 23  
 24  func (t *ServerTool) Info() agent.ToolInfo {
 25  	return agent.ToolInfo{
 26  		Name:        t.schema.Name,
 27  		Description: t.schema.Description,
 28  		Parameters:  t.schema.Parameters,
 29  	}
 30  }
 31  
 32  func (t *ServerTool) Run(ctx context.Context, argsJSON string) (agent.ToolResult, error) {
 33  	var args map[string]any
 34  	if argsJSON != "" && argsJSON != "{}" {
 35  		if err := json.Unmarshal([]byte(argsJSON), &args); err != nil {
 36  			return agent.ToolResult{
 37  				Content: fmt.Sprintf("invalid arguments: %v", err),
 38  				IsError: true,
 39  			}, nil
 40  		}
 41  	}
 42  	if args == nil {
 43  		args = map[string]any{}
 44  	}
 45  
 46  	resp, err := t.gateway.ExecuteTool(ctx, t.schema.Name, args, "")
 47  	if err != nil {
 48  		msg := err.Error()
 49  		prefix := classifyServerError(msg)
 50  		return agent.ToolResult{
 51  			Content: fmt.Sprintf("%sserver tool error: %v", prefix, err),
 52  			IsError: true,
 53  		}, nil
 54  	}
 55  
 56  	// Convert server-reported usage (xAI Grok tokens for x_search, SerpAPI
 57  	// queries for web_search, etc.) into an agent-level ToolUsage. Populated
 58  	// on the ToolResult so the audit logger can attribute cost per call; also
 59  	// emitted via context so the per-run usage accumulator picks it up.
 60  	// Server populates resp.Usage when the underlying provider returns billing
 61  	// info; older servers leave it nil and this is a no-op.
 62  	var toolUsage *agent.ToolUsage
 63  	if resp.Usage != nil {
 64  		u := resp.Usage
 65  		// The gateway currently returns a flat `tokens` count (synthetic for
 66  		// SERP tools, real input+output sum for x_search). If explicit
 67  		// input/output breakdowns are present, prefer them; else collapse
 68  		// `tokens` into TotalTokens so the accumulator still sees the volume.
 69  		totalTokens := u.TotalTokens
 70  		if totalTokens == 0 {
 71  			totalTokens = u.Tokens
 72  		}
 73  		if totalTokens == 0 {
 74  			totalTokens = u.InputTokens + u.OutputTokens
 75  		}
 76  		model := u.Model
 77  		if model == "" {
 78  			model = u.CostModel
 79  		}
 80  		toolUsage = &agent.ToolUsage{
 81  			Provider:     u.Provider,
 82  			Model:        model,
 83  			InputTokens:  u.InputTokens,
 84  			OutputTokens: u.OutputTokens,
 85  			TotalTokens:  totalTokens,
 86  			CostUSD:      u.CostUSD,
 87  			Units:        u.Units,
 88  			UnitType:     u.UnitType,
 89  		}
 90  		agent.EmitUsage(ctx, agent.TurnUsage{
 91  			InputTokens:  u.InputTokens,
 92  			OutputTokens: u.OutputTokens,
 93  			TotalTokens:  totalTokens,
 94  			CostUSD:      u.CostUSD,
 95  			// Gateway tool calls are not LLM calls from the driving model's
 96  			// perspective — leave LLMCalls=0 so session LLMCalls stays clean.
 97  			Model: model,
 98  		})
 99  	}
100  
101  	if resp.Error != nil && *resp.Error != "" {
102  		return agent.ToolResult{Content: *resp.Error, IsError: true, Usage: toolUsage}, nil
103  	}
104  
105  	if !resp.Success {
106  		return agent.ToolResult{Content: "tool execution failed", IsError: true, Usage: toolUsage}, nil
107  	}
108  
109  	// Prefer pre-formatted text from backend; fall back to raw JSON output
110  	if resp.Text != nil && *resp.Text != "" {
111  		return agent.ToolResult{Content: *resp.Text, Usage: toolUsage}, nil
112  	}
113  	if len(resp.Output) == 0 || string(resp.Output) == "null" {
114  		return agent.ToolResult{Content: "no output", Usage: toolUsage}, nil
115  	}
116  	return agent.ToolResult{Content: string(resp.Output), Usage: toolUsage}, nil
117  }
118  
119  // RequiresApproval returns false — the server enforces its own access control.
120  func (t *ServerTool) RequiresApproval() bool { return false }
121  
122  // classifyServerError returns the appropriate error prefix based on the error
123  // message, so the agent loop's error-handling instructions can guide the model
124  // to retry transient failures instead of fabricating explanations.
125  //
126  // Status-code markers (returned NNN) are checked before free-text transient
127  // keywords so that a 4xx response body mentioning "timeout" (e.g. validation
128  // "timeout must be <= 30") is not mis-tagged as transient and retried.
129  func classifyServerError(msg string) string {
130  	lower := strings.ToLower(msg)
131  	// Status-code classification first — the HTTP status is authoritative.
132  	if strings.Contains(lower, "returned 401") || strings.Contains(lower, "returned 403") {
133  		return "[permission error] "
134  	}
135  	if strings.Contains(lower, "returned 400") || strings.Contains(lower, "returned 422") {
136  		return "[validation error] "
137  	}
138  	if strings.Contains(lower, "returned 429") ||
139  		strings.Contains(lower, "returned 502") ||
140  		strings.Contains(lower, "returned 503") ||
141  		strings.Contains(lower, "returned 504") {
142  		return "[transient error] "
143  	}
144  	// Keyword fallback for network-layer failures that have no HTTP status
145  	// (connection refused/reset, DNS, timeouts before the server responded).
146  	for _, sig := range []string{
147  		"rate limit", "timeout", "timed out", "connection refused",
148  		"connection reset", "eof", "unavailable",
149  	} {
150  		if strings.Contains(lower, sig) {
151  			return "[transient error] "
152  		}
153  	}
154  	return ""
155  }
156  
157  // ToolSource implements agent.ToolSourcer for deterministic tool ordering.
158  func (t *ServerTool) ToolSource() agent.ToolSource { return agent.SourceGateway }