From a33d8ab949aac53575331ac37d18deebebee7959 Mon Sep 17 00:00:00 2001
From: precious112 <preciouskent8@gmail.com>
Date: Wed, 11 Mar 2026 08:58:20 +0100
Subject: [PATCH] made another fix for chat truncation

---
 .../versions/0014_message_tool_call_id.py     | 31 +++++++++++++++++++
 .../agent/src/argus_agent/agent/memory.py     |  4 ++-
 .../agent/src/argus_agent/queue/worker.py     |  1 +
 .../agent/src/argus_agent/storage/models.py   |  1 +
 4 files changed, 36 insertions(+), 1 deletion(-)
 create mode 100644 packages/agent/alembic/versions/0014_message_tool_call_id.py

diff --git a/packages/agent/alembic/versions/0014_message_tool_call_id.py b/packages/agent/alembic/versions/0014_message_tool_call_id.py
new file mode 100644
index 0000000..2b54f35
--- /dev/null
+++ b/packages/agent/alembic/versions/0014_message_tool_call_id.py
@@ -0,0 +1,31 @@
+"""Add tool_call_id column to messages table.
+
+Stores the LLM-assigned tool_call_id so multi-turn conversations
+can correctly link tool results to their parent assistant tool_calls.
+
+Revision ID: 0014
+Revises: 0013
+Create Date: 2026-03-11
+"""
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+from alembic import op
+
+revision: str = "0014"
+down_revision: str | None = "0013"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "messages",
+        sa.Column("tool_call_id", sa.String(64), server_default="", nullable=False),
+    )
+
+
+def downgrade() -> None:
+    op.drop_column("messages", "tool_call_id")
diff --git a/packages/agent/src/argus_agent/agent/memory.py b/packages/agent/src/argus_agent/agent/memory.py
index df19e0a..39df48d 100644
--- a/packages/agent/src/argus_agent/agent/memory.py
+++ b/packages/agent/src/argus_agent/agent/memory.py
@@ -48,6 +48,7 @@ async def persist_message(
         role: str,
         content: str = "",
         tool_calls: list[dict] | None = None,
+        tool_call_id: str = "",
         tool_result: dict | None = None,
         token_count: int = 0,
     ) -> str:
@@ -60,6 +61,7 @@ async def persist_message(
                 role=role,
                 content=content,
                 tool_calls=tool_calls,
+                tool_call_id=tool_call_id,
                 tool_result=tool_result,
                 token_count=token_count,
                 tenant_id=get_tenant_id(),
@@ -96,7 +98,7 @@ async def load_history(self) -> None:
                     self.messages.append(LLMMessage(
                         role="tool",
                         content=row.content,
-                        tool_call_id=row.id,
+                        tool_call_id=row.tool_call_id or row.id,
                         name="",
                     ))
 
diff --git a/packages/agent/src/argus_agent/queue/worker.py b/packages/agent/src/argus_agent/queue/worker.py
index 034dd9b..db0d69f 100644
--- a/packages/agent/src/argus_agent/queue/worker.py
+++ b/packages/agent/src/argus_agent/queue/worker.py
@@ -240,6 +240,7 @@ async def on_event(event_type: str, data: dict[str, Any]) -> None:
                             role=msg.role,
                             content=msg.content,
                             tool_calls=msg.tool_calls if msg.tool_calls else None,
+                            tool_call_id=msg.tool_call_id or "",
                             token_count=0,
                         )
                 except Exception:
diff --git a/packages/agent/src/argus_agent/storage/models.py b/packages/agent/src/argus_agent/storage/models.py
index 6a4fc0b..97767e4 100644
--- a/packages/agent/src/argus_agent/storage/models.py
+++ b/packages/agent/src/argus_agent/storage/models.py
@@ -40,6 +40,7 @@ class Message(Base):
     role: Mapped[str] = mapped_column(String(20))  # user, assistant, system, tool
     content: Mapped[str] = mapped_column(Text, default="")
     tool_calls: Mapped[dict | None] = mapped_column(JSON, nullable=True)
+    tool_call_id: Mapped[str] = mapped_column(String(64), default="")
     tool_result: Mapped[dict | None] = mapped_column(JSON, nullable=True)
     token_count: Mapped[int] = mapped_column(Integer, default=0)
     created_at: Mapped[datetime] = mapped_column(DateTime, default=_utcnow)