Python: route prompt-injection models through the AIPrompt concept

Expose system/user classification on AIPrompt and contribute role-sensitive framework sinks and data-extension sinks through that concept. Keep query customizations framework-independent without changing their results. Add a concept regression test covering all six SDK families, document the supported frameworks, and correct the release notes for the new queries and library API.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: BazookaMusic

docs/codeql/reusables/supported-frameworks.rst

@@ -281,6 +281,12 @@ and the CodeQL library pack ``codeql/python-all`` (`changelog <https://github.co
    libxml2, XML processing library
    lxml, XML processing library
    xmltodict, XML processing library
+   anthropic, AI model SDK
+   google-genai, AI model SDK
+   langchain, AI framework
+   openai, AI model SDK
+   openai-agents, AI agent framework
+   openrouter, AI model SDK
 
 
 Ruby built-in support

python/ql/lib/change-notes/2026-06-18-prompt-injection-models.md

@@ -1,4 +1,4 @@
 ---
-category: minorAnalysis
+category: feature
 ---
-* Added prompt-injection sink models (`system-prompt-injection` and `user-prompt-injection` kinds) for the `openai`, `agents`, `anthropic`, `google-genai`, `openrouter` and `langchain` frameworks.
+* Added the `AIPrompt` concept to model AI prompts and distinguish system-level prompts from user-level prompts. Added prompt-injection sink models (`system-prompt-injection` and `user-prompt-injection` kinds) for the `openai`, OpenAI Agents SDK (`agents`), `anthropic`, `google-genai`, `openrouter`, and `langchain` frameworks.

python/ql/lib/semmle/python/Concepts.qll

@@ -1804,6 +1804,12 @@ module Cryptography {
 class AIPrompt extends DataFlow::Node instanceof AIPrompt::Range {
   /** Gets an input that is used as AI prompt. */
   DataFlow::Node getAPrompt() { result = super.getAPrompt() }
+
+  /** Holds if this prompt supplies system-level instructions or assistant content. */
+  predicate isSystemPrompt() { super.isSystemPrompt() }
+
+  /** Holds if this prompt supplies user-level content. */
+  predicate isUserPrompt() { super.isUserPrompt() }
 }
 
 /** Provides a class for modeling new AI prompting mechanisms. */
@@ -1817,5 +1823,11 @@ module AIPrompt {
   abstract class Range extends DataFlow::Node {
     /** Gets an input that is used as AI prompt. */
     abstract DataFlow::Node getAPrompt();
+
+    /** Holds if this prompt supplies system-level instructions or assistant content. */
+    predicate isSystemPrompt() { none() }
+
+    /** Holds if this prompt supplies user-level content. */
+    predicate isUserPrompt() { none() }
   }
 }

python/ql/lib/semmle/python/Frameworks.qll

@@ -12,6 +12,7 @@ private import semmle.python.frameworks.Aiomysql
 private import semmle.python.frameworks.Aiopg
 private import semmle.python.frameworks.Aiosqlite
 private import semmle.python.frameworks.Airspeed
+private import semmle.python.frameworks.Anthropic
 private import semmle.python.frameworks.Anyio
 private import semmle.python.frameworks.Asyncpg
 private import semmle.python.frameworks.Baize
@@ -34,6 +35,7 @@ private import semmle.python.frameworks.Flask
 private import semmle.python.frameworks.FlaskAdmin
 private import semmle.python.frameworks.FlaskSqlAlchemy
 private import semmle.python.frameworks.Genshi
+private import semmle.python.frameworks.GoogleGenAI
 private import semmle.python.frameworks.Gradio
 private import semmle.python.frameworks.Hdbcli
 private import semmle.python.frameworks.Httpx
@@ -54,6 +56,8 @@ private import semmle.python.frameworks.Multidict
 private import semmle.python.frameworks.Mysql
 private import semmle.python.frameworks.MySQLdb
 private import semmle.python.frameworks.Numpy
+private import semmle.python.frameworks.OpenAI
+private import semmle.python.frameworks.OpenRouter
 private import semmle.python.frameworks.Opml
 private import semmle.python.frameworks.Oracledb
 private import semmle.python.frameworks.Pandas

python/ql/lib/semmle/python/frameworks/Anthropic.qll

@@ -11,6 +11,7 @@
 
 private import python
 private import semmle.python.ApiGraphs
+private import semmle.python.Concepts
 
 /** Provides classes modeling prompt-injection sinks of the `anthropic` package. */
 module Anthropic {
@@ -55,4 +56,14 @@ module Anthropic {
       result = msg.getSubscript("content")
     )
   }
+
+  private class MessagePrompt extends AIPrompt::Range {
+    MessagePrompt() { this = [getSystemOrAssistantPromptNode(), getUserPromptNode()].asSink() }
+
+    override DataFlow::Node getAPrompt() { result = this }
+
+    override predicate isSystemPrompt() { this = getSystemOrAssistantPromptNode().asSink() }
+
+    override predicate isUserPrompt() { this = getUserPromptNode().asSink() }
+  }
 }

python/ql/lib/semmle/python/frameworks/GoogleGenAI.qll

@@ -11,6 +11,7 @@
 
 private import python
 private import semmle.python.ApiGraphs
+private import semmle.python.Concepts
 
 /** Provides classes modeling prompt-injection sinks of the `google-genai` package. */
 module GoogleGenAI {
@@ -55,4 +56,14 @@ module GoogleGenAI {
       result = msg.getSubscript("parts").getASubscript().getSubscript("text")
     )
   }
+
+  private class ContentPrompt extends AIPrompt::Range {
+    ContentPrompt() { this = [getSystemOrAssistantPromptNode(), getUserPromptNode()].asSink() }
+
+    override DataFlow::Node getAPrompt() { result = this }
+
+    override predicate isSystemPrompt() { this = getSystemOrAssistantPromptNode().asSink() }
+
+    override predicate isUserPrompt() { this = getUserPromptNode().asSink() }
+  }
 }

python/ql/lib/semmle/python/frameworks/OpenAI.qll

@@ -14,6 +14,7 @@
 
 private import python
 private import semmle.python.ApiGraphs
+private import semmle.python.Concepts
 
 /** Holds if `msg` is a message dictionary with a privileged (system/developer/assistant) role. */
 private predicate isSystemOrDevMessage(API::Node msg) {
@@ -65,6 +66,16 @@ module AgentSdk {
       result = msg.getSubscript("content")
     )
   }
+
+  private class RunnerPrompt extends AIPrompt::Range {
+    RunnerPrompt() { this = [getSystemOrAssistantPromptNode(), getUserPromptNode()].asSink() }
+
+    override DataFlow::Node getAPrompt() { result = this }
+
+    override predicate isSystemPrompt() { this = getSystemOrAssistantPromptNode().asSink() }
+
+    override predicate isUserPrompt() { this = getUserPromptNode().asSink() }
+  }
 }
 
 /**
@@ -174,4 +185,14 @@ module OpenAI {
       result = realtimeItemContent(item)
     )
   }
+
+  private class ClientPrompt extends AIPrompt::Range {
+    ClientPrompt() { this = [getSystemOrAssistantPromptNode(), getUserPromptNode()].asSink() }
+
+    override DataFlow::Node getAPrompt() { result = this }
+
+    override predicate isSystemPrompt() { this = getSystemOrAssistantPromptNode().asSink() }
+
+    override predicate isUserPrompt() { this = getUserPromptNode().asSink() }
+  }
 }

python/ql/lib/semmle/python/frameworks/OpenRouter.qll

@@ -8,6 +8,7 @@
 
 private import python
 private import semmle.python.ApiGraphs
+private import semmle.python.Concepts
 
 /** Holds if `msg` is a message dictionary with a privileged (system/developer/assistant) role. */
 private predicate isSystemOrDevMessage(API::Node msg) {
@@ -57,4 +58,14 @@ module OpenRouter {
       result = messageContent(msg)
     )
   }
+
+  private class MessagePrompt extends AIPrompt::Range {
+    MessagePrompt() { this = [getSystemOrAssistantPromptNode(), getUserPromptNode()].asSink() }
+
+    override DataFlow::Node getAPrompt() { result = this }
+
+    override predicate isSystemPrompt() { this = getSystemOrAssistantPromptNode().asSink() }
+
+    override predicate isUserPrompt() { this = getUserPromptNode().asSink() }
+  }
 }

python/ql/lib/semmle/python/frameworks/data/ModelsAsData.qll

@@ -35,6 +35,22 @@ private class ThreatModelSourceFromDataExtension extends ThreatModelSource::Rang
   }
 }
 
+private class AIPromptFromDataExtension extends AIPrompt::Range {
+  AIPromptFromDataExtension() {
+    this = ModelOutput::getASinkNode(["system-prompt-injection", "user-prompt-injection"]).asSink()
+  }
+
+  override DataFlow::Node getAPrompt() { result = this }
+
+  override predicate isSystemPrompt() {
+    this = ModelOutput::getASinkNode("system-prompt-injection").asSink()
+  }
+
+  override predicate isUserPrompt() {
+    this = ModelOutput::getASinkNode("user-prompt-injection").asSink()
+  }
+}
+
 private class SummarizedCallableFromModel extends SummarizedCallable::Range {
   string type;
   string path;

python/ql/lib/semmle/python/security/dataflow/SystemPromptInjectionCustomizations.qll

@@ -9,11 +9,6 @@ private import semmle.python.Concepts
 private import semmle.python.ApiGraphs
 private import semmle.python.dataflow.new.RemoteFlowSources
 private import semmle.python.dataflow.new.BarrierGuards
-private import semmle.python.frameworks.data.ModelsAsData
-private import semmle.python.frameworks.OpenAI
-private import semmle.python.frameworks.Anthropic
-private import semmle.python.frameworks.GoogleGenAI
-private import semmle.python.frameworks.OpenRouter
 
 /**
  * Provides default sources, sinks and sanitizers for detecting
@@ -45,25 +40,7 @@ module SystemPromptInjection {
    * A prompt to an AI model, considered as a flow sink.
    */
   class AIPromptAsSink extends Sink {
-    AIPromptAsSink() { this = any(AIPrompt p).getAPrompt() }
-  }
-
-  private class SinkFromModel extends Sink {
-    SinkFromModel() { this = ModelOutput::getASinkNode("system-prompt-injection").asSink() }
-  }
-
-  private class PromptContentSink extends Sink {
-    PromptContentSink() {
-      this = OpenAI::getSystemOrAssistantPromptNode().asSink()
-      or
-      this = AgentSdk::getSystemOrAssistantPromptNode().asSink()
-      or
-      this = Anthropic::getSystemOrAssistantPromptNode().asSink()
-      or
-      this = GoogleGenAI::getSystemOrAssistantPromptNode().asSink()
-      or
-      this = OpenRouter::getSystemOrAssistantPromptNode().asSink()
-    }
+    AIPromptAsSink() { this = any(AIPrompt p | p.isSystemPrompt()).getAPrompt() }
   }
 
   /**