Cyrene_v2/internal/runtime/agent/step.go

package agent

import (
	"context"
	"fmt"
	"time"

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

// 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
}

// 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.
func (a *Agent) RunStep(ctx context.Context, sessionID string, messages []model.Message) (*StepResult, error) {
	req := a.buildRequest(messages)

	if a.guards != nil {
		verdict, err := a.guards.CheckInput(ctx, &req)
		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 {
		verdict, err := a.guards.CheckOutput(ctx, &resp)
		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
}

// 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
}