Make It Bulletproof

async Task<string> CallMcpToolWithRetry(
    string toolName, string argsJson,
    int maxRetries = 3, double backoff = 1.0)
{
    for (var attempt = 0; attempt < maxRetries; attempt++)
    {
        try
        {
            return await CallMcpTool(toolName, argsJson);
        }
        catch (TaskCanceledException)
        {
            var wait = TimeSpan.FromSeconds(backoff * Math.Pow(2, attempt));
            _logger.LogWarning("Tool {Tool} timed out (attempt {Attempt}/{Max})",
                toolName, attempt + 1, maxRetries);
            await Task.Delay(wait);
        }
        catch (HttpRequestException ex)
        {
            _logger.LogError(ex, "Tool {Tool} HTTP error", toolName);
            return $"Error: tool {toolName} failed: {ex.StatusCode}";
        }
    }
    return $"Error: tool {toolName} timed out after {maxRetries} retries";
}

export async function callMcpToolWithRetry(
  name: string,
  args: Record<string, any>,
  maxRetries = 3,
  backoff = 1000,
): Promise<string> {
  for (let attempt = 0; attempt < maxRetries; attempt++) {
    try {
      return await callMcpTool(name, args);
    } catch (err: any) {
      if (err.name === 'AbortError' || err.code === 'ETIMEDOUT') {
        const wait = backoff * Math.pow(2, attempt);
        console.warn(`Tool ${name} timed out (attempt ${attempt + 1}), retrying in ${wait}ms`);
        await new Promise(r => setTimeout(r, wait));
        continue;
      }
      console.error(`Tool ${name} failed:`, err.message);
      return `Error: tool ${name} failed: ${err.message}`;
    }
  }
  return `Error: tool ${name} timed out after ${maxRetries} retries`;
}

func callMcpToolWithRetry(name string, args map[string]interface{}, maxRetries int) (string, error) {
    backoff := time.Second
    for attempt := 0; attempt < maxRetries; attempt++ {
        result, err := callMcpTool(name, args)
        if err == nil {
            return result, nil
        }
        if os.IsTimeout(err) {
            wait := backoff * time.Duration(1<<attempt)
            log.Printf("Tool %s timed out (attempt %d/%d), retrying in %v",
                name, attempt+1, maxRetries, wait)
            time.Sleep(wait)
            continue
        }
        log.Printf("Tool %s failed: %v", name, err)
        return fmt.Sprintf("Error: tool %s failed: %v", name, err), nil
    }
    return fmt.Sprintf("Error: tool %s timed out after %d retries", name, maxRetries), nil
}

// Retry the transport, never the tool logic. A tool that answered with
// isError=true gave you a real answer - retrying it just burns time and money.

private static final int MAX_ATTEMPTS = 3;

private <T> T postWithRetry(JsonRpcRequest request, Class<T> responseType) {
    RuntimeException last = null;

    for (var attempt = 1; attempt <= MAX_ATTEMPTS; attempt++) {
        try {
            var response = http.post()
                    .uri(baseUrl + "/message")
                    .body(request)
                    .retrieve()
                    .body(responseType);

            if (response == null) {
                throw new IllegalStateException("Empty response from MCP server");
            }
            return response;
        } catch (RestClientException | IllegalStateException e) {
            last = e instanceof RuntimeException re ? re : new IllegalStateException(e);
            if (attempt == MAX_ATTEMPTS) {
                break;
            }

            // Exponential backoff: 500ms, then 1s. Keep it short - the user is
            // waiting on a chat response, not a batch job.
            var backoff = Duration.ofMillis(500L * (1L << (attempt - 1)));
            log.warn("MCP call failed (attempt {}/{}): {}. Retrying in {}ms",
                    attempt, MAX_ATTEMPTS, e.getMessage(), backoff.toMillis());
            try {
                Thread.sleep(backoff);
            } catch (InterruptedException ie) {
                Thread.currentThread().interrupt();
                throw new IllegalStateException("Interrupted while retrying MCP call", ie);
            }
        }
    }

    throw new IllegalStateException("MCP call failed after %d attempts".formatted(MAX_ATTEMPTS), last);
}

private readonly ILogger<McpAgent> _logger;

async Task<string> ProcessQuery(string query, string sessionId)
{
    var sw = Stopwatch.StartNew();
    _logger.LogInformation("Query start: {Session} {Query}",
        sessionId, query[..Math.Min(200, query.Length)]);

    var iteration = 0;
    while (true)
    {
        iteration++;
        var result = await openai.CompleteChatAsync(messages, options);

        if (result.Value.FinishReason != ChatFinishReason.ToolCalls)
        {
            _logger.LogInformation(
                "Query complete: {Session} iterations={Iter} duration={Ms}ms",
                sessionId, iteration, sw.ElapsedMilliseconds);
            return result.Value.Content[0].Text;
        }

        foreach (var tc in result.Value.ToolCalls)
            _logger.LogInformation("Tool call: {Tool} {Args} iter={Iter}",
                tc.FunctionName, tc.FunctionArguments, iteration);
    }
}

function log(event: string, data: Record<string, any>) {
  console.log(JSON.stringify({ ts: new Date().toISOString(), event, ...data }));
}

async function processQuery(query: string, sessionId = 'default'): Promise<string> {
  const start = Date.now();
  log('query_start', { sessionId, query: query.slice(0, 200) });

  let iteration = 0;
  while (true) {
    iteration++;
    const response = await openai.chat.completions.create({
      model: MODEL, messages, tools: tools.length ? tools : undefined,
    });

    const message = response.choices[0].message;
    if (!message.tool_calls?.length) {
      log('query_complete', {
        sessionId, iterations: iteration,
        durationMs: Date.now() - start,
        tokens: response.usage?.total_tokens,
      });
      return message.content ?? '';
    }

    for (const tc of message.tool_calls)
      log('tool_call', { tool: tc.function.name, args: tc.function.arguments, iteration });
  }
}

func processQuery(query, sessionID string) (string, error) {
    start := time.Now()
    log.Printf(`{"event":"query_start","session":"%s","query":"%s"}`,
        sessionID, truncate(query, 200))

    iteration := 0
    for {
        iteration++
        resp, err := aiClient.CreateChatCompletion(ctx, req)
        if err != nil {
            return "", fmt.Errorf("LLM error: %w", err)
        }

        choice := resp.Choices[0]
        if len(choice.Message.ToolCalls) == 0 {
            log.Printf(`{"event":"query_complete","session":"%s","iterations":%d,"duration_ms":%d}`,
                sessionID, iteration, time.Since(start).Milliseconds())
            return choice.Message.Content, nil
        }

        for _, tc := range choice.Message.ToolCalls {
            log.Printf(`{"event":"tool_call","tool":"%s","iteration":%d}`,
                tc.Function.Name, iteration)
        }
    }
}

// Spring Boot 4 emits structured JSON with one property. No logging library,
// no custom appender - and every line becomes queryable in your log backend.

// application.properties
logging.structured.format.console=ecs

// Then MDC keys become top-level fields on every line in that scope. Put the
// things you will actually filter on there: which tool, which round, how long.

try (var ignored = MDC.putCloseable("tool", name);
     var ignored2 = MDC.putCloseable("round", String.valueOf(round))) {

    var started = System.nanoTime();
    var result = mcp.callTool(name, arguments);
    var elapsedMs = Duration.ofNanos(System.nanoTime() - started).toMillis();

    // Structured arguments, not string concatenation. "{}" placeholders keep the
    // message template stable so you can group on it.
    log.info("tool call finished in {}ms, {} chars returned", elapsedMs, result.length());
    return result;
}

public class ContextManager
{
    private const int MaxTokens = 8000;

    // Rough estimate: ~4 chars per token for English
    static int EstimateTokens(IList<ChatMessage> messages)
    {
        return messages.Sum(m =>
            4 + (m.ToString()?.Length ?? 0) / 4);
    }

    public static void TrimMessages(List<ChatMessage> messages)
    {
        while (EstimateTokens(messages) > MaxTokens && messages.Count > 2)
            messages.RemoveAt(1); // Keep system prompt
    }
}

const MAX_CONTEXT_TOKENS = 8000;

// Rough estimate: ~4 chars per token for English text
function estimateTokens(messages: any[]): number {
  return messages.reduce((sum, m) => {
    const content = typeof m.content === 'string' ? m.content : JSON.stringify(m);
    return sum + 4 + Math.ceil(content.length / 4);
  }, 0);
}

export function trimMessages(messages: any[], max = MAX_CONTEXT_TOKENS): any[] {
  while (estimateTokens(messages) > max && messages.length > 2) {
    messages.splice(1, 1); // Keep system prompt at index 0
  }
  return messages;
}

const maxContextTokens = 8000

// Rough estimate: ~4 chars per token
func estimateTokens(messages []openai.ChatCompletionMessage) int {
    total := 0
    for _, m := range messages {
        total += 4 + len(m.Content)/4
    }
    return total
}

func trimMessages(messages *[]openai.ChatCompletionMessage) {
    for estimateTokens(*messages) > maxContextTokens && len(*messages) > 2 {
        // Remove second message, keep system prompt at index 0
        *messages = append((*messages)[:1], (*messages)[2:]...)
    }
}

// Every round appends an assistant message and one tool message per call, so a
// long conversation grows the prompt on every turn. Trim the middle, never the
// ends: the system prompt sets the rules and the newest turns carry the intent.

private static final int MAX_HISTORY_MESSAGES = 20;

private List<ChatCompletionMessageParam> trimHistory(List<ChatCompletionMessageParam> messages) {
    if (messages.size() <= MAX_HISTORY_MESSAGES) {
        return messages;
    }

    var trimmed = new ArrayList<ChatCompletionMessageParam>();
    trimmed.add(messages.getFirst());                       // keep the system prompt

    // Drop from the front of the tail, but never orphan a tool message: a tool
    // message whose assistant tool_calls entry was dropped is a 400 from the API.
    var tail = messages.subList(messages.size() - (MAX_HISTORY_MESSAGES - 1), messages.size());
    var firstKeepable = 0;
    while (firstKeepable < tail.size() && tail.get(firstKeepable).isTool()) {
        firstKeepable++;
    }

    trimmed.addAll(tail.subList(firstKeepable, tail.size()));
    log.info("Trimmed history from {} to {} messages", messages.size(), trimmed.size());
    return trimmed;
}

const int MaxToolResultChars = 4000;

static string TruncateResult(string result)
{
    if (result.Length <= MaxToolResultChars)
        return result;
    return result[..MaxToolResultChars] +
        $"\n\n[Truncated: showing {MaxToolResultChars}/{result.Length} chars]";
}

const MAX_TOOL_RESULT_CHARS = 4000;

export function truncateResult(result: string, max = MAX_TOOL_RESULT_CHARS): string {
  if (result.length <= max) return result;
  return result.slice(0, max) +
    `\n\n[Truncated: showing ${max}/${result.length} chars]`;
}

const maxToolResultChars = 4000

func truncateResult(result string) string {
    if len(result) <= maxToolResultChars {
        return result
    }
    return result[:maxToolResultChars] +
        fmt.Sprintf("\n\n[Truncated: showing %d/%d chars]", maxToolResultChars, len(result))
}

// A tool that returns a 200 KB payload will happily eat your whole context
// window. Cut it, and say so in the text - a silent truncation makes the model
// confidently summarise data it never saw.

private static final int MAX_RESULT_CHARS = 4000;

private static String truncate(String result) {
    if (result.length() <= MAX_RESULT_CHARS) {
        return result;
    }

    var kept = result.substring(0, MAX_RESULT_CHARS);
    var dropped = result.length() - MAX_RESULT_CHARS;

    return kept + "%n%n[truncated: %d of %d characters omitted]".formatted(dropped, result.length());
}

public class RateLimiter
{
    private readonly int _max;
    private readonly TimeSpan _window;
    private readonly ConcurrentDictionary<string, List<DateTime>> _requests = new();

    public RateLimiter(int maxRequests = 10, int windowSeconds = 60)
    {
        _max = maxRequests;
        _window = TimeSpan.FromSeconds(windowSeconds);
    }

    public bool Check(string key)
    {
        var now = DateTime.UtcNow;
        var timestamps = _requests.GetOrAdd(key, _ => new List<DateTime>());
        lock (timestamps)
        {
            timestamps.RemoveAll(t => now - t > _window);
            if (timestamps.Count >= _max) return false;
            timestamps.Add(now);
            return true;
        }
    }
}

export class RateLimiter {
  private requests = new Map<string, number[]>();

  constructor(
    private maxRequests = 10,
    private windowMs = 60_000,
  ) {}

  check(key: string): boolean {
    const now = Date.now();
    const timestamps = (this.requests.get(key) ?? [])
      .filter(t => now - t < this.windowMs);

    if (timestamps.length >= this.maxRequests) return false;
    timestamps.push(now);
    this.requests.set(key, timestamps);
    return true;
  }
}

type RateLimiter struct {
    mu       sync.Mutex
    max      int
    window   time.Duration
    requests map[string][]time.Time
}

func NewRateLimiter(max int, window time.Duration) *RateLimiter {
    return &RateLimiter{max: max, window: window, requests: make(map[string][]time.Time)}
}

func (rl *RateLimiter) Check(key string) bool {
    rl.mu.Lock()
    defer rl.mu.Unlock()
    now := time.Now()
    // Remove expired
    valid := rl.requests[key][:0]
    for _, t := range rl.requests[key] {
        if now.Sub(t) < rl.window {
            valid = append(valid, t)
        }
    }
    if len(valid) >= rl.max {
        rl.requests[key] = valid
        return false
    }
    rl.requests[key] = append(valid, now)
    return true
}

package no.javazone.agent;

import java.time.Duration;
import java.util.concurrent.Semaphore;
import org.springframework.stereotype.Component;

// A token bucket in twenty lines. Refills on a scheduled tick, blocks callers
// when empty. Guards the OpenAI quota, which is the resource you actually run
// out of first.
@Component
class RateLimiter {

    private static final int CAPACITY = 10;

    private final Semaphore tokens = new Semaphore(CAPACITY);

    // Give back one token per second, but never exceed the bucket size -
    // otherwise an idle agent accumulates a burst it can spend all at once.
    @Scheduled(fixedRate = 1000)
    void refill() {
        if (tokens.availablePermits() < CAPACITY) {
            tokens.release();
        }
    }

    void acquire() {
        try {
            if (!tokens.tryAcquire(Duration.ofSeconds(30))) {
                throw new IllegalStateException("Rate limit exceeded, try again shortly");
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new IllegalStateException("Interrupted while waiting for rate limit", e);
        }
    }
}

const int MaxIterations = 10;

for (var i = 0; i < MaxIterations; i++)
{
    var result = await openai.CompleteChatAsync(messages, options);
    if (result.Value.FinishReason != ChatFinishReason.ToolCalls)
        return result.Value.Content[0].Text;

    ContextManager.TrimMessages(messages);
    // ... execute tool calls
}

_logger.LogWarning("Hit max iterations for query");
return "Sorry, I couldn't complete this request.";

const MAX_ITERATIONS = 10;

for (let i = 0; i < MAX_ITERATIONS; i++) {
  const response = await openai.chat.completions.create({
    model: MODEL, messages, tools: tools.length ? tools : undefined,
  });

  const message = response.choices[0].message;
  if (!message.tool_calls?.length) return message.content ?? '';

  trimMessages(messages);
  // ... execute tool calls
}

console.warn('Hit max iterations');
return "Sorry, I couldn't complete this request.";

const maxIterations = 10

for i := 0; i < maxIterations; i++ {
    resp, _ := aiClient.CreateChatCompletion(ctx, req)
    choice := resp.Choices[0]
    if len(choice.Message.ToolCalls) == 0 {
        return choice.Message.Content, nil
    }
    trimMessages(&messages)
    // ... execute tool calls
}

log.Println("Hit max iterations")
return "Sorry, I couldn't complete this request.", nil

// Three guards, and each one stops a different failure. The round cap stops a
// model that keeps calling the same tool. The duplicate check stops the tighter
// version of that, where it calls the same tool with the same arguments. The
// deadline stops a slow upstream from holding a request open forever.

private static final int MAX_ROUNDS = 10;
private static final Duration DEADLINE = Duration.ofSeconds(60);

String processQuery(String query) {
    var deadline = java.time.Instant.now().plus(DEADLINE);
    var seen = new HashSet<String>();

    for (var round = 0; round < MAX_ROUNDS; round++) {
        if (java.time.Instant.now().isAfter(deadline)) {
            throw new IllegalStateException("Agent exceeded its %ds deadline".formatted(DEADLINE.toSeconds()));
        }

        rateLimiter.acquire();
        var completion = openai.chat().completions().create(builder.build());
        var message = completion.choices().getFirst().message();
        var toolCalls = message.toolCalls().orElse(List.of());

        if (toolCalls.isEmpty()) {
            return message.content().orElse("(No response generated)");
        }

        builder.addMessage(message);

        for (var toolCall : toolCalls) {
            var functionCall = toolCall.function().orElse(null);
            if (functionCall == null) {
                continue;
            }

            // Same tool, same arguments, second time: tell the model instead of
            // paying for the call again. It usually moves on when told.
            var fingerprint = functionCall.name() + functionCall.arguments();
            if (!seen.add(fingerprint)) {
                builder.addMessage(ChatCompletionToolMessageParam.builder()
                        .toolCallId(toolCall.id())
                        .content("This tool was already called with these arguments. Use the earlier result.")
                        .build());
                continue;
            }

            builder.addMessage(runTool(functionCall));
        }
    }

    throw new IllegalStateException("Agent loop exceeded " + MAX_ROUNDS + " iterations");
}