r/src/context_manager.c

// retoor <retoor@molodetz.nl>
#include "context_manager.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#define MIN_KEEP_CHARS 500
#define TRUNCATE_MARKER "\n\n[... content truncated for context management ...]\n\n"
static const char *get_message_role(struct json_object *msg) {
    struct json_object *role_obj;
    if (json_object_object_get_ex(msg, "role", &role_obj)) {
        return json_object_get_string(role_obj);
    }
    // Tool results don't have a 'role' field in some implementations, 
    // they have 'tool_call_id'.
    if (json_object_object_get_ex(msg, "tool_call_id", &role_obj)) {
        return "tool";
    }
    return "";
}
static bool has_tool_calls(struct json_object *msg) {
    struct json_object *tool_calls;
    if (json_object_object_get_ex(msg, "tool_calls", &tool_calls)) {
        return json_object_array_length(tool_calls) > 0;
    }
    return false;
}
static size_t get_message_content_len(struct json_object *msg) {
    struct json_object *content_obj;
    const char *content = NULL;
    if (json_object_object_get_ex(msg, "content", &content_obj)) {
        content = json_object_get_string(content_obj);
    } else if (json_object_object_get_ex(msg, "tool_result", &content_obj)) {
        content = json_object_get_string(content_obj);
    }
    return content ? strlen(content) : 0;
}
static size_t calculate_total_size(messages_handle msgs) {
    size_t total = 0;
    int count = messages_count(msgs);
    for (int i = 0; i < count; i++) {
        total += get_message_content_len(messages_get_object(msgs, i));
    }
    return total;
}
static r_status_t perform_truncate(messages_handle msgs, int index, double ratio) {
    struct json_object *msg = messages_get_object(msgs, index);
    if (!msg) return R_ERROR_NOT_FOUND;
    struct json_object *content_obj;
    const char *content = NULL;
    bool is_tool_result = false;
    if (json_object_object_get_ex(msg, "content", &content_obj)) {
        content = json_object_get_string(content_obj);
    } else if (json_object_object_get_ex(msg, "tool_result", &content_obj)) {
        content = json_object_get_string(content_obj);
        is_tool_result = true;
    }
    if (!content) return R_SUCCESS;
    size_t len = strlen(content);
    size_t target_len = (size_t)(len * ratio);
    if (target_len < MIN_KEEP_CHARS * 2) target_len = MIN_KEEP_CHARS * 2;
    if (target_len >= len) return R_SUCCESS;
    size_t keep_each = target_len / 2;
    
    size_t marker_len = strlen(TRUNCATE_MARKER);
    char *new_content = malloc(keep_each * 2 + marker_len + 1);
    if (!new_content) return R_ERROR_OUT_OF_MEMORY;
    memcpy(new_content, content, keep_each);
    memcpy(new_content + keep_each, TRUNCATE_MARKER, marker_len);
    size_t offset = keep_each + marker_len;
    memcpy(new_content + offset, content + len - keep_each, keep_each);
    new_content[offset + keep_each] = '\0';
    struct json_object *new_msg = json_tokener_parse(json_object_to_json_string(msg));
    if (is_tool_result) {
        json_object_object_add(new_msg, "tool_result", json_object_new_string(new_content));
    } else {
        json_object_object_add(new_msg, "content", json_object_new_string(new_content));
    }
    free(new_content);
    return messages_replace_at(msgs, index, new_msg);
}
r_status_t context_manager_shrink(messages_handle msgs) {
    if (!msgs) return R_ERROR_INVALID_ARG;
    int count = messages_count(msgs);
    if (count <= 2) return R_ERROR_API_ERROR;
    size_t initial_size = calculate_total_size(msgs);
    size_t target_size = (size_t)(initial_size * 0.5);
    if (target_size < 20000) target_size = 20000;
    fprintf(stderr, "  \033[2m-> Context overflow (%zu chars). Middle-out shrinking to %zu...\033[0m\n", 
            initial_size, target_size);
    // Strategy 1: Truncate very large messages first (safe, doesn't break sequence)
    for (int i = 0; i < messages_count(msgs); i++) {
        struct json_object *msg = messages_get_object(msgs, i);
        if (get_message_content_len(msg) > 50000) {
            perform_truncate(msgs, i, 0.2);
        }
    }
    // Strategy 2: Remove messages from the middle until size is within target
    // We keep: 
    // - System message (usually index 0)
    // - Most recent 4 messages (usually current task context)
    while (calculate_total_size(msgs) > target_size && messages_count(msgs) > 6) {
        int middle_idx = 1; // Start after system
        struct json_object *msg = messages_get_object(msgs, middle_idx);
        const char *role = get_message_role(msg);
        int remove_count = 1;
        if (strcmp(role, "assistant") == 0 && has_tool_calls(msg)) {
            // Must also remove the following tool results to maintain sequence
            int search_idx = middle_idx + 1;
            while (search_idx < messages_count(msgs) - 4) {
                struct json_object *next_msg = messages_get_object(msgs, search_idx);
                if (strcmp(get_message_role(next_msg), "tool") == 0) {
                    remove_count++;
                    search_idx++;
                } else {
                    break;
                }
            }
        }
        // Ensure we don't eat into the "recent" buffer
        if (middle_idx + remove_count > messages_count(msgs) - 4) {
            break;
        }
        messages_remove_range(msgs, middle_idx, remove_count);
    }
    size_t final_size = calculate_total_size(msgs);
    fprintf(stderr, "  \033[2m-> Context shrunk to %zu chars. Remaining messages: %d\033[0m\n", 
            final_size, messages_count(msgs));
    return R_SUCCESS;
}