Cyrene_v2/internal/runtime/agent/step.go

package agent

import (
	"context"
	"fmt"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/guardrail"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/tool"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/trace"
)

// Stop reasons carried by StepResult.StopReason, consumed by the loop (M2) to
// decide how to continue after a step.
const (
	// StopReasonStop is a normal final answer with no tool calls.
	StopReasonStop = "stop"
	// StopReasonToolCalls means the assistant requested tool calls.
	StopReasonToolCalls = "tool_calls"
	// StopReasonGuardrailInput means the input guardrail blocked the request.
	StopReasonGuardrailInput = "guardrail_input"
	// StopReasonGuardrailOutput means the output guardrail blocked the response.
	StopReasonGuardrailOutput = "guardrail_output"
)

// ToolCall is a single function call the assistant requested, decoded from a
// step. Args holds the tool arguments as a JSON object.
type ToolCall struct {
	// ID is the provider tool_use/tool_call id, so the loop can feed the result
	// back as a native tool_result correlated to the exact call (Anthropic
	// tool_use_id / OpenAI tool_call_id). Empty for the legacy JSON-envelope
	// convention, where the provider emitted no id.
	ID   string         `json:"id,omitempty"`
	Name string         `json:"name"`
	Args map[string]any `json:"args"`
}

// StepResult is the outcome of a single agent step.
type StepResult struct {
	// Final is true when the step produced a final answer (or was intercepted by
	// a guardrail) and no tool execution is required.
	Final bool
	// Text is the assistant's final text, or the guardrail's block reason when
	// Final is true.
	Text string
	// ToolCalls is non-empty when Final is false, listing the tool calls to run.
	ToolCalls []ToolCall
	// StopReason classifies how the step ended (see the StopReason* consts).
	StopReason string
	// ModelStopReason is the provider's raw stop/finish reason for the model
	// call (e.g. Anthropic max_tokens / end_turn / tool_use), so a trace can
	// distinguish a token-budget truncation from a normal end. Empty when the
	// step never reached the model (guardrail denial).
	ModelStopReason string
}

// 护栏 span 名与本包记录事件时使用的类型：与循环侧（loop/spanlog.go）
// 同名，使面板无论走流式路径还是 agent 路径都看到同一类 guardrail 节点
// 与同一条日志事件（裁决 11.4.2，design/04 §9「逐层记到 guardrail」）。
const (
	guardrailSpanInput  = "guardrail.input"
	guardrailSpanOutput = "guardrail.output"
	guardrailEventDone  = "guardrail.done"
)

// RunStep executes one step of the agent loop. It assembles a model.Request from
// the agent's instructions prepended as a leading system message plus the
// conversation messages, runs input guardrails, calls the model adapter, runs
// output guardrails and decodes any tool calls. sessionID identifies the
// conversation for downstream tracing/telemetry and is currently carried without
// a sink; it is reserved for trace/persist (M2/M9).
//
// A final text answer yields Final=true with StopReason "stop"; a tool request
// yields Final=false with the decoded ToolCalls; a guardrail denial yields
// Final=true carrying the block reason. A nil adapter or a failed guardrail or
// model call returns a wrapped error.
//
// 两侧检查都包在自己的 guardrail span 里（裁决 11.4.2）：本包不持有 Tracer，
// 因此用 trace.ChildSpan 在当前 span（循环的 step span）之下开一个子 span，
// 未开追踪时它是空操作。**本路径保持串行内联**：输入侧在模型调用之前拦下，
// 输出侧紧接模型返回裁决（裁决 13.4.1 的「输出侧异步」由流式主路径
// loop.streamStep 承担——agent.RunStep 只服务静音/非流式内部流程，不落在
// 主会话关键路径上，且它的裁决是 StepResult 的一部分）。
func (a *Agent) RunStep(ctx context.Context, sessionID string, messages []model.Message) (*StepResult, error) {
	req := a.buildRequest(messages)

	if a.guards != nil {
		gctx, gspan := trace.ChildSpan(ctx, guardrailSpanInput, trace.SpanKindGuardrail)
		verdict, err := a.guards.CheckInput(gctx, &req)
		finishGuardrailSpan(gspan, guardrailSpanInput, verdict, err)
		if err != nil {
			return nil, fmt.Errorf("agent: check input: %w", err)
		}
		if verdict != nil && !verdict.Allow {
			return &StepResult{Final: true, Text: verdict.Reason, StopReason: StopReasonGuardrailInput}, nil
		}
	}

	if a.adapter == nil {
		return nil, fmt.Errorf("agent: run step: %w", ErrNoAdapter)
	}

	start := time.Now()
	resp, err := a.adapter.Chat(ctx, req, model.Options{})
	if err != nil {
		return nil, fmt.Errorf("agent: chat: %w", err)
	}
	if a.recorder != nil {
		a.recorder(model.CallRecord{
			Model:   resp.Model,
			Purpose: req.Purpose,
			Tokens:  resp.TotalTokens,
			TookMs:  time.Since(start).Milliseconds(),
		})
	}

	if a.guards != nil {
		gctx, gspan := trace.ChildSpan(ctx, guardrailSpanOutput, trace.SpanKindGuardrail)
		verdict, err := a.guards.CheckOutput(gctx, &resp)
		finishGuardrailSpan(gspan, guardrailSpanOutput, verdict, err)
		if err != nil {
			return nil, fmt.Errorf("agent: check output: %w", err)
		}
		if verdict != nil && !verdict.Allow {
			return &StepResult{Final: true, Text: verdict.Reason, StopReason: StopReasonGuardrailOutput}, nil
		}
	}

	if calls, ok := ExtractToolCalls(resp); ok {
		return &StepResult{Final: false, ToolCalls: calls, StopReason: StopReasonToolCalls, ModelStopReason: resp.StopReason}, nil
	}
	return &StepResult{Final: true, Text: resp.Content, StopReason: StopReasonStop, ModelStopReason: resp.StopReason}, nil
}

// finishGuardrailSpan 把一次护栏检查的结论落到其 span 上（裁决 11.4.2）：
// props 只放检查名、裁决与耗时（全部在图的 safePropKeys 白名单内，出图时仍会
// 再清洗一遍），命中拦截或标记时额外记一条脱敏后的 Reason 事件。拦截的 span
// 终态为 error，面板上显红。span 为 nil（未开追踪）时是空操作。
func finishGuardrailSpan(span *trace.Span, action string, verdict *guardrail.Verdict, err error) {
	if span == nil {
		return
	}
	allow := err == nil && (verdict == nil || verdict.Allow)
	reason := ""
	if verdict != nil {
		reason = verdict.Reason
	}
	if err != nil {
		reason = err.Error()
	}
	_ = span.SetProp("action", action)
	_ = span.SetProp("allow", allow)
	_ = span.SetProp("elapsed_ms", trace.ElapsedMS(span))
	if err != nil {
		_ = span.SetProp("error", err.Error())
	}
	if !allow || reason != "" {
		_ = span.Record(trace.Event{Type: guardrailEventDone, Props: map[string]any{
			"action": action,
			"allow":  allow,
			"reason": truncateRunes(reason, maxGuardrailReasonRunes),
		}})
	}
	if allow {
		trace.FinishChild(span, trace.SpanStatusOK, nil)
		return
	}
	trace.FinishChild(span, trace.SpanStatusError, nil)
}

// maxGuardrailReasonRunes 限制护栏 Reason 进入 span 事件前的长度：Reason 是
// 面向用户的拦截说明，本就不含密钥（guardrail 包契约），这里只防它变成一大段。
const maxGuardrailReasonRunes = 200

// truncateRunes 把 s 截到最多 max 个 rune，超出时补省略号；max 非正时原样返回。
func truncateRunes(s string, max int) string {
	if max <= 0 {
		return s
	}
	r := []rune(s)
	if len(r) <= max {
		return s
	}
	return string(r[:max]) + "…"
}

// ExtractToolCalls decodes tool calls from a model.Response. It prefers the
// provider-native ToolCalls the adapters populate (Anthropic tool_use / OpenAI
// tool_calls), which carry the provider id needed to correlate the result and
// are the only shape a real native-tool provider emits. It then falls back to
// two legacy conventions used by prompt-engineered providers:
//
//  1. The normalized assistant-message placeholder: when resp.Message carries a
//     ToolCallName it is a single call with no argument payload.
//  2. A JSON envelope in resp.Content of the form
//     {"tool_calls":[{"name":"...","args":{...}}]}.
//
// It reports ok=true (with a non-empty slice) when a call is decoded, and
// ok=false otherwise.
func ExtractToolCalls(resp model.Response) ([]ToolCall, bool) {
	if len(resp.ToolCalls) > 0 {
		return nativeToolCalls(resp.ToolCalls), true
	}
	if resp.Message.Role == model.RoleAssistant && resp.Message.ToolCallName != "" {
		return []ToolCall{{Name: resp.Message.ToolCallName, Args: map[string]any{}}}, true
	}
	var payload struct {
		ToolCalls []ToolCall `json:"tool_calls"`
	}
	// Lenient decode repairs common model-emitted malformations (code fences,
	// trailing commas) so a slightly broken tool-call envelope does not silently
	// turn into a text answer; an unrecoverable payload still yields ok=false and
	// the round proceeds as text (the loop never crashes on a bad call).
	if err := tool.LenientJSONUnmarshal([]byte(resp.Content), &payload); err != nil {
		return nil, false
	}
	if len(payload.ToolCalls) == 0 {
		return nil, false
	}
	return payload.ToolCalls, true
}

// nativeToolCalls converts the adapter-decoded provider tool calls into the
// agent's ToolCall, preserving the provider id and normalizing a nil Args to an
// empty object so downstream fingerprinting and invocation never see nil.
func nativeToolCalls(calls []model.ToolCall) []ToolCall {
	out := make([]ToolCall, 0, len(calls))
	for _, c := range calls {
		args := c.Args
		if args == nil {
			args = map[string]any{}
		}
		out = append(out, ToolCall{ID: c.ID, Name: c.Name, Args: args})
	}
	return out
}